Tackling Water Pollution using YOLO-NAS and W&B

Introduction

Table of Contents

The Problem

The Dataset

Dataset Loading and Version Control

# Download Dataset for the first time
from roboflow import Roboflow
rf = Roboflow(api_key=os.getenv("YOUR_ROBOFLOW_API_KEY"))
project = rf.workspace("easyhyeon").project("trash-sea")
dataset = project.version(10).download("yolov5")

# Upload as a WandB Artifact
with wandb.init(project=PROJECT_NAME, entity=ENTITY, resume="allow", save_code=True)
    artifact = wandb.Artifact('trash-sea-10', type='dataset')
    artifact.add_dir(DATASET_PATH)
    wandb.log_artifact(artifact)

# Get Artifact
with wandb.init(project=PROJECT_NAME, entity=ENTITY):
    artifact = run.use_artifact(
	'ml-colabs/fconn-yolo-nas/trash-sea-10:latest',
	type='dataset'
    )
    artifact_dir = artifact.download()

Dataset Registration

# Define our variables for the dataset
dataset_params = {
    'data_dir':DATASET_PATH,
    'train_images_dir':'train/images', # The relative path for training images
    'train_labels_dir':'train/labels', # The relative path for training annotations
    'val_images_dir':'valid/images', # The relative path for validation images
    'val_labels_dir':'valid/labels', # The relative path for validation annotations
    'test_images_dir':'test/images', # The relative path for test images
    'test_labels_dir':'test/labels', # The relative path for test annotations
    'classes': ["Buoy", "Can", "Paper", "Plastic Bag", "Plastic Bottle"] # Classes
}

train_data = coco_detection_yolo_format_train(
    dataset_params={
        'data_dir': dataset_params['data_dir'],
        'images_dir': dataset_params['train_images_dir'],
        'labels_dir': dataset_params['train_labels_dir'],
        'classes': dataset_params['classes'],
    },
    dataloader_params={
        'batch_size':config["batch_size"], # Set batch_size according to experiments
        'num_workers':2 
        # This variable defines the number of workers that your CPU will initialize
	# and use for loading data. If your runtime has a high-CPU configuration,
        # feel free to use the num_workers argument as 2 or more. Else, reduce it
	# to 0 to have the main-process perform dataset fetching into the RAM.
    }
)

val_data = coco_detection_yolo_format_val(
    dataset_params={
        'data_dir': dataset_params['data_dir'],
        'images_dir': dataset_params['val_images_dir'],
        'labels_dir': dataset_params['val_labels_dir'],
        'classes': dataset_params['classes'],
    },
    dataloader_params={
        'batch_size':config["batch_size"],
        'num_workers':2
    }
)

test_data = coco_detection_yolo_format_val(
    dataset_params={
        'data_dir': dataset_params['data_dir'],
        'images_dir': dataset_params['test_images_dir'],
        'labels_dir': dataset_params['test_labels_dir'],
        'classes': dataset_params['classes'],
    },
    dataloader_params={
        'batch_size':config["batch_size"],
        'num_workers':2
    }
)

Exploratory Data Analysis



Defining our learning objectives

We have 5 specific classes (Buoy, Can, Paper, Plastic Bag, and Plastic Bottle) wherein the number of objects varies significantly.
The dataset shows a deficiency of the Can and Plastic Bag classes but has an abundance of the Buoy and Plastic Bottle objects. This indicates a class imbalance that should be corrected.
The bottom-right area of the frame seems to have an abundance of objects to be detected and is populated by 3 different classes. This may be a concern wherein translation-invariance will play an important role.
Plastic Bottles and Paper have been captured from various angles and at various zoom levels, thus contributing to a large variation within the bounding box sizes for this class.
Buoys are normally captured from afar, contributing to a low variation within the size of the bounding boxes for this class.
All classes have a significant number of outliers when it comes to bounding box size.

The Model

YOLO Class of Models

YOLO-NAS Class of Models



Use of advanced regimes during the pre-training stage, such as knowledge distillation, post-training quantization (through the introduction of RepVGG blocks into the model architecture), quantization-aware training, and the introduction of a Distribution Focal Loss term.
Quantization-compatible approach wherein the Neural Architecture Search algorithm adaptively quantized layers, skipping those which may have a significant impact on the network's accuracy/latency trade-off. As a result, the network supports high compression up to INT8 quantization while suffering a minuscule drop in precision.

YOLO-NAS-S: With 19M params, this is the smallest YOLO-NAS variant. Deci.ai reports an impressive performance of 47.5 mAP on the COCO 2017 Validation dataset. This is supposed to be the most nimble and performant model in its class, owing to its small size.
YOLO-NAS-M: This model has up to 51.1M parameters, and is the medium-sized version to provide a balance between size and performance. It reports a performance of up to 51.55 mAP on the Validation set.
YOLO-NAS-L: This model is the largest of its class, at 66.9M params. It is the most performant model, giving up to 52.22 mAP on the Validation set.

Experimentations and Actionable Insights from our Baselines

Baseline Experiments

Translating Insights into Performance



Final Model Training

Inference on Test Data - Performance with Test Distributions

Conclusion

Load our data present in the YOLO dataset format using SuperGradients.
Instantiate and use the latest YOLO-NAS model from Deci.ai.
Perform dataset analysis using Weave by Weights & Biases.
Create baselines for each model variant available.
Initiate a full WandB Sweep for 200 runs and get the best set of hyperparameters.
Train a final model for a longer time and more number of epochs.
Test the model on unseen data and gather results.