.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "intermediate/torchvision_tutorial.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note Click :ref:`here ` to download the full example code .. rst-class:: sphx-glr-example-title .. _sphx_glr_intermediate_torchvision_tutorial.py: TorchVision Object Detection Finetuning Tutorial ==================================================== .. GENERATED FROM PYTHON SOURCE LINES 8-107 For this tutorial, we will be finetuning a pre-trained `Mask R-CNN `_ model on the `Penn-Fudan Database for Pedestrian Detection and Segmentation `_. It contains 170 images with 345 instances of pedestrians, and we will use it to illustrate how to use the new features in torchvision in order to train an object detection and instance segmentation model on a custom dataset. .. note :: This tutorial works only with torchvision version >=0.16 or nightly. If you're using torchvision<=0.15, please follow `this tutorial instead `_. Defining the Dataset -------------------- The reference scripts for training object detection, instance segmentation and person keypoint detection allows for easily supporting adding new custom datasets. The dataset should inherit from the standard :class:`torch.utils.data.Dataset` class, and implement ``__len__`` and ``__getitem__``. The only specificity that we require is that the dataset ``__getitem__`` should return a tuple: - image: :class:`torchvision.tv_tensors.Image` of shape ``[3, H, W]``, a pure tensor, or a PIL Image of size ``(H, W)`` - target: a dict containing the following fields - ``boxes``, :class:`torchvision.tv_tensors.BoundingBoxes` of shape ``[N, 4]``: the coordinates of the ``N`` bounding boxes in ``[x0, y0, x1, y1]`` format, ranging from ``0`` to ``W`` and ``0`` to ``H`` - ``labels``, integer :class:`torch.Tensor` of shape ``[N]``: the label for each bounding box. ``0`` represents always the background class. - ``image_id``, int: an image identifier. It should be unique between all the images in the dataset, and is used during evaluation - ``area``, float :class:`torch.Tensor` of shape ``[N]``: the area of the bounding box. This is used during evaluation with the COCO metric, to separate the metric scores between small, medium and large boxes. - ``iscrowd``, uint8 :class:`torch.Tensor` of shape ``[N]``: instances with ``iscrowd=True`` will be ignored during evaluation. - (optionally) ``masks``, :class:`torchvision.tv_tensors.Mask` of shape ``[N, H, W]``: the segmentation masks for each one of the objects If your dataset is compliant with above requirements then it will work for both training and evaluation codes from the reference script. Evaluation code will use scripts from ``pycocotools`` which can be installed with ``pip install pycocotools``. .. note :: For Windows, please install ``pycocotools`` from `gautamchitnis `_ with command ``pip install git+https://github.com/gautamchitnis/cocoapi.git@cocodataset-master#subdirectory=PythonAPI`` One note on the ``labels``. The model considers class ``0`` as background. If your dataset does not contain the background class, you should not have ``0`` in your ``labels``. For example, assuming you have just two classes, *cat* and *dog*, you can define ``1`` (not ``0``) to represent *cats* and ``2`` to represent *dogs*. So, for instance, if one of the images has both classes, your ``labels`` tensor should look like ``[1, 2]``. Additionally, if you want to use aspect ratio grouping during training (so that each batch only contains images with similar aspect ratios), then it is recommended to also implement a ``get_height_and_width`` method, which returns the height and the width of the image. If this method is not provided, we query all elements of the dataset via ``__getitem__`` , which loads the image in memory and is slower than if a custom method is provided. Writing a custom dataset for PennFudan ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Let’s write a dataset for the PennFudan dataset. First, let's download the dataset and extract the `zip file `_: .. code:: python wget https://www.cis.upenn.edu/~jshi/ped_html/PennFudanPed.zip -P data cd data && unzip PennFudanPed.zip We have the following folder structure: :: PennFudanPed/ PedMasks/ FudanPed00001_mask.png FudanPed00002_mask.png FudanPed00003_mask.png FudanPed00004_mask.png ... PNGImages/ FudanPed00001.png FudanPed00002.png FudanPed00003.png FudanPed00004.png Here is one example of a pair of images and segmentation masks .. GENERATED FROM PYTHON SOURCE LINES 108-124 .. code-block:: default import matplotlib.pyplot as plt from torchvision.io import read_image image = read_image("data/PennFudanPed/PNGImages/FudanPed00046.png") mask = read_image("data/PennFudanPed/PedMasks/FudanPed00046_mask.png") plt.figure(figsize=(16, 8)) plt.subplot(121) plt.title("Image") plt.imshow(image.permute(1, 2, 0)) plt.subplot(122) plt.title("Mask") plt.imshow(mask.permute(1, 2, 0)) .. image-sg:: /intermediate/images/sphx_glr_torchvision_tutorial_001.png :alt: Image, Mask :srcset: /intermediate/images/sphx_glr_torchvision_tutorial_001.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none .. GENERATED FROM PYTHON SOURCE LINES 125-137 So each image has a corresponding segmentation mask, where each color correspond to a different instance. Let’s write a :class:`torch.utils.data.Dataset` class for this dataset. In the code below, we are wrapping images, bounding boxes and masks into :class:`torchvision.tv_tensors.TVTensor` classes so that we will be able to apply torchvision built-in transformations (`new Transforms API `_) for the given object detection and segmentation task. Namely, image tensors will be wrapped by :class:`torchvision.tv_tensors.Image`, bounding boxes into :class:`torchvision.tv_tensors.BoundingBoxes` and masks into :class:`torchvision.tv_tensors.Mask`. As :class:`torchvision.tv_tensors.TVTensor` are :class:`torch.Tensor` subclasses, wrapped objects are also tensors and inherit the plain :class:`torch.Tensor` API. For more information about torchvision ``tv_tensors`` see `this documentation `_. .. GENERATED FROM PYTHON SOURCE LINES 137-202 .. code-block:: default import os import torch from torchvision.io import read_image from torchvision.ops.boxes import masks_to_boxes from torchvision import tv_tensors from torchvision.transforms.v2 import functional as F class PennFudanDataset(torch.utils.data.Dataset): def __init__(self, root, transforms): self.root = root self.transforms = transforms # load all image files, sorting them to # ensure that they are aligned self.imgs = list(sorted(os.listdir(os.path.join(root, "PNGImages")))) self.masks = list(sorted(os.listdir(os.path.join(root, "PedMasks")))) def __getitem__(self, idx): # load images and masks img_path = os.path.join(self.root, "PNGImages", self.imgs[idx]) mask_path = os.path.join(self.root, "PedMasks", self.masks[idx]) img = read_image(img_path) mask = read_image(mask_path) # instances are encoded as different colors obj_ids = torch.unique(mask) # first id is the background, so remove it obj_ids = obj_ids[1:] num_objs = len(obj_ids) # split the color-encoded mask into a set # of binary masks masks = (mask == obj_ids[:, None, None]).to(dtype=torch.uint8) # get bounding box coordinates for each mask boxes = masks_to_boxes(masks) # there is only one class labels = torch.ones((num_objs,), dtype=torch.int64) image_id = idx area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0]) # suppose all instances are not crowd iscrowd = torch.zeros((num_objs,), dtype=torch.int64) # Wrap sample and targets into torchvision tv_tensors: img = tv_tensors.Image(img) target = {} target["boxes"] = tv_tensors.BoundingBoxes(boxes, format="XYXY", canvas_size=F.get_size(img)) target["masks"] = tv_tensors.Mask(masks) target["labels"] = labels target["image_id"] = image_id target["area"] = area target["iscrowd"] = iscrowd if self.transforms is not None: img, target = self.transforms(img, target) return img, target def __len__(self): return len(self.imgs) .. GENERATED FROM PYTHON SOURCE LINES 203-238 That’s all for the dataset. Now let’s define a model that can perform predictions on this dataset. Defining your model ------------------- In this tutorial, we will be using `Mask R-CNN `_, which is based on top of `Faster R-CNN `_. Faster R-CNN is a model that predicts both bounding boxes and class scores for potential objects in the image. .. image:: ../../_static/img/tv_tutorial/tv_image03.png Mask R-CNN adds an extra branch into Faster R-CNN, which also predicts segmentation masks for each instance. .. image:: ../../_static/img/tv_tutorial/tv_image04.png There are two common situations where one might want to modify one of the available models in TorchVision Model Zoo. The first is when we want to start from a pre-trained model, and just finetune the last layer. The other is when we want to replace the backbone of the model with a different one (for faster predictions, for example). Let’s go see how we would do one or another in the following sections. 1 - Finetuning from a pretrained model ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Let’s suppose that you want to start from a model pre-trained on COCO and want to finetune it for your particular classes. Here is a possible way of doing it: .. GENERATED FROM PYTHON SOURCE LINES 238-254 .. code-block:: default import torchvision from torchvision.models.detection.faster_rcnn import FastRCNNPredictor # load a model pre-trained on COCO model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights="DEFAULT") # replace the classifier with a new one, that has # num_classes which is user-defined num_classes = 2 # 1 class (person) + background # get number of input features for the classifier in_features = model.roi_heads.box_predictor.cls_score.in_features # replace the pre-trained head with a new one model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes) .. rst-class:: sphx-glr-script-out .. code-block:: none Downloading: "https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth" to /var/lib/ci-user/.cache/torch/hub/checkpoints/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth 0%| | 0.00/160M [00:00`_ to easily write data augmentation pipelines for Object Detection and Segmentation tasks. Let’s write some helper functions for data augmentation / transformation: .. GENERATED FROM PYTHON SOURCE LINES 365-377 .. code-block:: default from torchvision.transforms import v2 as T def get_transform(train): transforms = [] if train: transforms.append(T.RandomHorizontalFlip(0.5)) transforms.append(T.ToDtype(torch.float, scale=True)) transforms.append(T.ToPureTensor()) return T.Compose(transforms) .. GENERATED FROM PYTHON SOURCE LINES 378-383 Testing ``forward()`` method (Optional) --------------------------------------- Before iterating over the dataset, it's good to see what the model expects during training and inference time on sample data. .. GENERATED FROM PYTHON SOURCE LINES 383-408 .. code-block:: default import utils model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights="DEFAULT") dataset = PennFudanDataset('data/PennFudanPed', get_transform(train=True)) data_loader = torch.utils.data.DataLoader( dataset, batch_size=2, shuffle=True, collate_fn=utils.collate_fn ) # For Training images, targets = next(iter(data_loader)) images = list(image for image in images) targets = [{k: v for k, v in t.items()} for t in targets] output = model(images, targets) # Returns losses and detections print(output) # For inference model.eval() x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)] predictions = model(x) # Returns predictions print(predictions[0]) .. rst-class:: sphx-glr-script-out .. code-block:: none {'loss_classifier': tensor(0.0798, grad_fn=), 'loss_box_reg': tensor(0.0284, grad_fn=), 'loss_objectness': tensor(0.0186, grad_fn=), 'loss_rpn_box_reg': tensor(0.0034, grad_fn=)} {'boxes': tensor([], size=(0, 4), grad_fn=), 'labels': tensor([], dtype=torch.int64), 'scores': tensor([], grad_fn=)} .. GENERATED FROM PYTHON SOURCE LINES 409-411 Let’s now write the main function which performs the training and the validation: .. GENERATED FROM PYTHON SOURCE LINES 411-481 .. code-block:: default from engine import train_one_epoch, evaluate # train on the GPU or on the CPU, if a GPU is not available device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') # our dataset has two classes only - background and person num_classes = 2 # use our dataset and defined transformations dataset = PennFudanDataset('data/PennFudanPed', get_transform(train=True)) dataset_test = PennFudanDataset('data/PennFudanPed', get_transform(train=False)) # split the dataset in train and test set indices = torch.randperm(len(dataset)).tolist() dataset = torch.utils.data.Subset(dataset, indices[:-50]) dataset_test = torch.utils.data.Subset(dataset_test, indices[-50:]) # define training and validation data loaders data_loader = torch.utils.data.DataLoader( dataset, batch_size=2, shuffle=True, collate_fn=utils.collate_fn ) data_loader_test = torch.utils.data.DataLoader( dataset_test, batch_size=1, shuffle=False, collate_fn=utils.collate_fn ) # get the model using our helper function model = get_model_instance_segmentation(num_classes) # move model to the right device model.to(device) # construct an optimizer params = [p for p in model.parameters() if p.requires_grad] optimizer = torch.optim.SGD( params, lr=0.005, momentum=0.9, weight_decay=0.0005 ) # and a learning rate scheduler lr_scheduler = torch.optim.lr_scheduler.StepLR( optimizer, step_size=3, gamma=0.1 ) # let's train it just for 2 epochs num_epochs = 2 for epoch in range(num_epochs): # train for one epoch, printing every 10 iterations train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10) # update the learning rate lr_scheduler.step() # evaluate on the test dataset evaluate(model, data_loader_test, device=device) print("That's it!") .. rst-class:: sphx-glr-script-out .. code-block:: none Downloading: "https://download.pytorch.org/models/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth" to /var/lib/ci-user/.cache/torch/hub/checkpoints/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth 0%| | 0.00/170M [00:00 50, and a mask mAP of 65. But what do the predictions look like? Let’s take one image in the dataset and verify .. GENERATED FROM PYTHON SOURCE LINES 488-518 .. code-block:: default import matplotlib.pyplot as plt from torchvision.utils import draw_bounding_boxes, draw_segmentation_masks image = read_image("data/PennFudanPed/PNGImages/FudanPed00046.png") eval_transform = get_transform(train=False) model.eval() with torch.no_grad(): x = eval_transform(image) # convert RGBA -> RGB and move to device x = x[:3, ...].to(device) predictions = model([x, ]) pred = predictions[0] image = (255.0 * (image - image.min()) / (image.max() - image.min())).to(torch.uint8) image = image[:3, ...] pred_labels = [f"pedestrian: {score:.3f}" for label, score in zip(pred["labels"], pred["scores"])] pred_boxes = pred["boxes"].long() output_image = draw_bounding_boxes(image, pred_boxes, pred_labels, colors="red") masks = (pred["masks"] > 0.7).squeeze(1) output_image = draw_segmentation_masks(output_image, masks, alpha=0.5, colors="blue") plt.figure(figsize=(12, 12)) plt.imshow(output_image.permute(1, 2, 0)) .. image-sg:: /intermediate/images/sphx_glr_torchvision_tutorial_002.png :alt: torchvision tutorial :srcset: /intermediate/images/sphx_glr_torchvision_tutorial_002.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none .. GENERATED FROM PYTHON SOURCE LINES 519-535 The results look good! Wrapping up ----------- In this tutorial, you have learned how to create your own training pipeline for object detection models on a custom dataset. For that, you wrote a :class:`torch.utils.data.Dataset` class that returns the images and the ground truth boxes and segmentation masks. You also leveraged a Mask R-CNN model pre-trained on COCO train2017 in order to perform transfer learning on this new dataset. For a more complete example, which includes multi-machine / multi-GPU training, check ``references/detection/train.py``, which is present in the torchvision repository. .. rst-class:: sphx-glr-timing **Total running time of the script:** ( 0 minutes 49.667 seconds) .. _sphx_glr_download_intermediate_torchvision_tutorial.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: torchvision_tutorial.py ` .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: torchvision_tutorial.ipynb ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_