Skip to content

Open In Colab

IceVision meets W&B

IceVision + W&B = Agnostic Object Detection Framework with Outstanding Experiments Tracking

IceVision fully supports W&B by providing a one-liner API that enables users to track their trained models and display both the predicted and ground truth bounding boxes.

W&B makes visualizing and tracking different models performance a highly enjoyable task. Indeed, we are able to monitor the performance of several EfficientDet backbones by changing few lines of code and obtaining very intuitive and easy-to-interpret figures that highlights both the similarities and differences between the different backbones.

For more information check the Report. Note, however, that the report refers to an older version of IceVision. This tutorial is updated for IceVision 0.7.

This tutorial emphasizes the additional work required to integrated W&B. If you are new to IceVision, then we suggest that you look at the getting started with object detection tutorial.

In this tutorial, we are using the fastai library training loop, the efficientdet object detection model, and a sample dataset with images of objects that you might find in a fridge. Following the usual practice with IceVision, you can use W&B with other training loops, model libraries, models and backbones. The W&B specific lines below would not need to be changed.

Install IceVision and IceData

If on Colab run the following cell, else check the installation instructions

# IceVision - IceData - MMDetection - YOLO v5 Installation
#!wget https://raw.githubusercontent.com/airctic/icevision/master/install_colab.sh
#!chmod +x install_colab.sh && ./install_colab.sh

Imports

from icevision.all import *
from fastai.callback.wandb import *
from fastai.callback.tracker import SaveModelCallback

Load the Fridge Objects dataset

The fridge Objects dataset is tiny dataset that contains 134 images of 4 classes: - can, - carton, - milk bottle, - water bottle.

IceVision provides very handy methods such as loading a dataset, parsing annotations, and more.

# Dataset
url = "https://cvbp-secondary.z19.web.core.windows.net/datasets/object_detection/odFridgeObjects.zip"
dest_dir = "fridge"
data_dir = icedata.load_data(url, dest_dir, force_download=True)
# Parser
class_map = ClassMap(["milk_bottle", "carton", "can", "water_bottle"])
parser = parsers.VOCBBoxParser(annotations_dir=data_dir / "odFridgeObjects/annotations",
                     images_dir=data_dir / "odFridgeObjects/images",
                     class_map=class_map)
# Records
train_records, valid_records = parser.parse(show_pbar=False)

Train and Validation Datasets

# Transforms
image_size = 384

train_tfms = tfms.A.Adapter([*tfms.A.aug_tfms(size=image_size, presize=512), tfms.A.Normalize()])
valid_tfms = tfms.A.Adapter([*tfms.A.resize_and_pad(image_size), tfms.A.Normalize()])
# Datasets
train_ds = Dataset(train_records, train_tfms)
valid_ds = Dataset(valid_records, valid_tfms)

Create the model

In IceVision, we need to select the model type and backbone. For this tutorial, we are selecting efficientdet and the tf_lite0 backbone. Some models require additional information, such as the image_size.

# Library and model selection

model_type = models.ross.efficientdet
backbone = model_type.backbones.tf_lite0(pretrained=True)
# The efficientdet model requires an img_size parameter
extra_args = {'img_size' : image_size}

model = efficientdet.model(backbone=backbone, num_classes=len(class_map), **extra_args) 

Create the dataloaders

The dataloaders differ somewhat across the model_types, so creating them comes after selecting the model type.

# DataLoaders
train_dl = model_type.train_dl(train_ds, batch_size=16, num_workers=4, shuffle=True)
valid_dl = model_type.valid_dl(valid_ds, batch_size=16, num_workers=4, shuffle=False)

Training

Intialize W&B

At this point, we initialize W&B. This works in the jupyter notebook, but it is more typical run W&B from within a programme. This is partly because it enables you to track the progress of your training jobs from a custom dashboard from your browser, tablet, or phone. The full interface also makes it easy to compare multiple training runs, which can be very powerful when combined with IceVision. You can easily see which model is best suited to your problem.

Initializing is a single line from the W&B library.

wandb.init(project="icevision-fridge", name="efficientdet_tf_lite0", reinit=True)

Create the learner

This tutorial is usingfastai, but IceVision lets you us other frameworks such as pytorch-lightning.

In order to use W&B within fastai, you need to specify the WandbCallback, which results in logging the metrics as well as other key parameters, as well as the SaveModelCallback, which enables W&B to log the models. Logging the model is very powerful, as it ensures that you have a copy of the best version of the model as you train. If you are using W&B on-line, however, it causes your model to be transferred to the W&B database as well as saved in a local wandb directory.

learn = model_type.fastai.learner(dls=[train_dl, valid_dl], model=model, 
                                  metrics=[COCOMetric(metric_type=COCOMetricType.bbox)], 
                                  cbs=[WandbCallback(), SaveModelCallback()])

Train

In this case, we use the fit_one_cycle training method from fastai, which uses a specific policy for adjusting the learning rate. This model is likely to take around 2-10 seconds per epoch, depending on your hardware. Training for 30 epochs on this small dataset typically reaches a level around 0.8 (COCOMetric), which is sufficient for our demonstration purposes and saves some time.

learn.fit_one_cycle(30, 1e-2)

Show results

We can now look athe results of the training in the notebook.

model_type.show_results(model, valid_ds)

Get predictions

Let's get the list of predictions from our model. We do this by creating an infer_dl - a dataloader used for inference and then getting predictions from the data loader.

Please note the keep_images=True. By default, the predictions include scores, labels, and bounding boxes. In our case, we want to keep the images so that we log them to W&B.

infer_dl = model_type.infer_dl(valid_ds, batch_size=8)
preds = model_type.predict_from_dl(model=model, infer_dl=infer_dl, keep_images=True)

Log results to W&B

Now comes the most important bit of this tutorial - actually logging the predictions to W&B. This takes one line specific to icevision and a second line to send the information to W&B.

# Create wandb_images for each prediction
wandb_images = wandb_img_preds(preds, add_ground_truth=True) 

# Log the wandb_images to wandb
wandb.log({"Predicted images": wandb_images})

After logging and finishing the training, it is good to mark the run as completed. This can take a few seconds, as we wait for the W&B processes to transfer data and finalize logging.

# optional: mark the run as completed
wandb.join()

Happy Learning!

If you need any assistance, feel free to join our forum.