torchvision_finetuning_instance_segmentation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DfPPQ6ztJhv4"
   },
   "source": [
    "# TorchVision Instance Segmentation Finetuning Tutorial\n",
    "\n",
    "For this tutorial, we will be finetuning a pre-trained [Mask R-CNN](https://arxiv.org/abs/1703.06870) model in the [*Penn-Fudan Database for Pedestrian Detection and Segmentation*](https://www.cis.upenn.edu/~jshi/ped_html/). 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 instance segmentation model on a custom dataset.\n",
    "\n",
    "First, we need to install `pycocotools`. This library will be used for computing the evaluation metrics following the COCO metric for intersection over union."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "DBIoe_tHTQgV",
    "outputId": "a8aa5ac8-fe40-46a2-d0af-106d7da6e149"
   },
   "outputs": [],
   "source": [
    "! pip install cython\n",
    "# Install pycocotools, the version by default in Colab\n",
    "# has a bug fixed in https://github.com/cocodataset/cocoapi/pull/354\n",
    "! pip install -U 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5Sd4jlGp2eLm"
   },
   "source": [
    "## Defining the Dataset\n",
    "\n",
    "The [torchvision reference scripts for training object detection, instance segmentation and person keypoint detection](https://github.com/pytorch/vision/tree/v0.3.0/references/detection) allows for easily supporting adding new custom datasets.\n",
    "The dataset should inherit from the standard `torch.utils.data.Dataset` class, and implement `__len__` and `__getitem__`.\n",
    "\n",
    "The only specificity that we require is that the dataset `__getitem__` should return:\n",
    "\n",
    "* image: a PIL Image of size (H, W)\n",
    "* target: a dict containing the following fields\n",
    "    * `boxes` (`FloatTensor[N, 4]`): the coordinates of the `N` bounding boxes in `[x0, y0, x1, y1]` format, ranging from `0` to `W` and `0` to `H`\n",
    "    * `labels` (`Int64Tensor[N]`): the label for each bounding box\n",
    "    * `image_id` (`Int64Tensor[1]`): an image identifier. It should be unique between all the images in the dataset, and is used during evaluation\n",
    "    * `area` (`Tensor[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.\n",
    "    * `iscrowd` (`UInt8Tensor[N]`): instances with `iscrowd=True` will be ignored during evaluation.\n",
    "    * (optionally) `masks` (`UInt8Tensor[N, H, W]`): The segmentation masks for each one of the objects\n",
    "    * (optionally) `keypoints` (`FloatTensor[N, K, 3]`): For each one of the `N` objects, it contains the `K` keypoints in `[x, y, visibility]` format, defining the object. `visibility=0` means that the keypoint is not visible. Note that for data augmentation, the notion of flipping a keypoint is dependent on the data representation, and you should probably adapt `references/detection/transforms.py` for your new keypoint representation\n",
    "\n",
    "If your model returns the above methods, they will make it work for both training and evaluation, and will use the evaluation scripts from pycocotools.\n",
    "\n",
    "\n",
    "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].\n",
    "\n",
    "Additionally, if you want to use aspect ratio grouping during training (so that each batch only contains images with similar aspect ratio), 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.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bX0rqK-A3Nbl"
   },
   "source": [
    "### Writing a custom dataset for Penn-Fudan\n",
    "\n",
    "Let's write a dataset for the Penn-Fudan dataset.\n",
    "\n",
    "First, let's download and extract the data, present in a zip file at https://www.cis.upenn.edu/~jshi/ped_html/PennFudanPed.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext ersa.client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run_remote wget https://www.cis.upenn.edu/~jshi/ped_html/PennFudanPed.zip\n",
    "%run_remote unzip PennFudanPed.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "_t4TBwhHTdkd",
    "outputId": "5b1f86ef-e9af-4625-ca43-92c21daf2186"
   },
   "outputs": [],
   "source": [
    "# download the Penn-Fudan dataset\n",
    "! wget https://www.cis.upenn.edu/~jshi/ped_html/PennFudanPed.zip\n",
    "# extract it in the current folder\n",
    "! unzip PennFudanPed.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WfwuU-jI3j93"
   },
   "source": [
    "Let's have a look at the dataset and how it is layed down.\n",
    "\n",
    "The data is structured as follows\n",
    "```\n",
    "PennFudanPed/\n",
    "  PedMasks/\n",
    "    FudanPed00001_mask.png\n",
    "    FudanPed00002_mask.png\n",
    "    FudanPed00003_mask.png\n",
    "    FudanPed00004_mask.png\n",
    "    ...\n",
    "  PNGImages/\n",
    "    FudanPed00001.png\n",
    "    FudanPed00002.png\n",
    "    FudanPed00003.png\n",
    "    FudanPed00004.png\n",
    "```\n",
    "\n",
    "Here is one example of an image in the dataset, with its corresponding instance segmentation mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 553
    },
    "id": "LDjuVFgexFfh",
    "outputId": "f2e0c5c5-a8d0-4999-c94d-1660f153cd00"
   },
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "\n",
    "Image.open(\"PennFudanPed/PNGImages/FudanPed00001.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 553
    },
    "id": "cFHKCvCTxiff",
    "outputId": "f0309090-7704-463f-c4e5-50faceb36dbd"
   },
   "outputs": [],
   "source": [
    "mask = Image.open(\"PennFudanPed/PedMasks/FudanPed00001_mask.png\")\n",
    "# each mask instance has a different color, from zero to N, where\n",
    "# N is the number of instances. In order to make visualization easier,\n",
    "# let's adda color palette to the mask.\n",
    "mask.putpalette(\n",
    "    [\n",
    "        0,\n",
    "        0,\n",
    "        0,  # black background\n",
    "        255,\n",
    "        0,\n",
    "        0,  # index 1 is red\n",
    "        255,\n",
    "        255,\n",
    "        0,  # index 2 is yellow\n",
    "        255,\n",
    "        153,\n",
    "        0,  # index 3 is orange\n",
    "    ]\n",
    ")\n",
    "mask"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "C9Ee5NV54Dmj"
   },
   "source": [
    "So each image has a corresponding segmentation mask, where each color correspond to a different instance. Let's write a `torch.utils.data.Dataset` class for this dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mTgWtixZTs3X"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.utils.data\n",
    "from PIL import Image\n",
    "\n",
    "\n",
    "class PennFudanDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, root, transforms=None):\n",
    "        self.root = root\n",
    "        self.transforms = transforms\n",
    "        # load all image files, sorting them to\n",
    "        # ensure that they are aligned\n",
    "        self.imgs = list(sorted(os.listdir(os.path.join(root, \"PNGImages\"))))\n",
    "        self.masks = list(sorted(os.listdir(os.path.join(root, \"PedMasks\"))))\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        # load images ad masks\n",
    "        img_path = os.path.join(self.root, \"PNGImages\", self.imgs[idx])\n",
    "        mask_path = os.path.join(self.root, \"PedMasks\", self.masks[idx])\n",
    "        img = Image.open(img_path).convert(\"RGB\")\n",
    "        # note that we haven't converted the mask to RGB,\n",
    "        # because each color corresponds to a different instance\n",
    "        # with 0 being background\n",
    "        mask = Image.open(mask_path)\n",
    "\n",
    "        mask = np.array(mask)\n",
    "        # instances are encoded as different colors\n",
    "        obj_ids = np.unique(mask)\n",
    "        # first id is the background, so remove it\n",
    "        obj_ids = obj_ids[1:]\n",
    "\n",
    "        # split the color-encoded mask into a set\n",
    "        # of binary masks\n",
    "        masks = mask == obj_ids[:, None, None]\n",
    "\n",
    "        # get bounding box coordinates for each mask\n",
    "        num_objs = len(obj_ids)\n",
    "        boxes = []\n",
    "        for i in range(num_objs):\n",
    "            pos = np.where(masks[i])\n",
    "            xmin = np.min(pos[1])\n",
    "            xmax = np.max(pos[1])\n",
    "            ymin = np.min(pos[0])\n",
    "            ymax = np.max(pos[0])\n",
    "            boxes.append([xmin, ymin, xmax, ymax])\n",
    "\n",
    "        boxes = torch.as_tensor(boxes, dtype=torch.float32)\n",
    "        # there is only one class\n",
    "        labels = torch.ones((num_objs,), dtype=torch.int64)\n",
    "        masks = torch.as_tensor(masks, dtype=torch.uint8)\n",
    "\n",
    "        image_id = torch.tensor([idx])\n",
    "        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])\n",
    "        # suppose all instances are not crowd\n",
    "        iscrowd = torch.zeros((num_objs,), dtype=torch.int64)\n",
    "\n",
    "        target = {}\n",
    "        target[\"boxes\"] = boxes\n",
    "        target[\"labels\"] = labels\n",
    "        target[\"masks\"] = masks\n",
    "        target[\"image_id\"] = image_id\n",
    "        target[\"area\"] = area\n",
    "        target[\"iscrowd\"] = iscrowd\n",
    "\n",
    "        if self.transforms is not None:\n",
    "            img, target = self.transforms(img, target)\n",
    "\n",
    "        return img, target\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.imgs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "J6f3ZOTJ4Km9"
   },
   "source": [
    "That's all for the dataset. Let's see how the outputs are structured for this dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ZEARO4B_ye0s",
    "outputId": "061226b6-45c4-4b51-b650-05687cd8f67a"
   },
   "outputs": [],
   "source": [
    "dataset = PennFudanDataset(\"PennFudanPed/\")\n",
    "dataset[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lWOhcsir9Ahx"
   },
   "source": [
    "So we can see that by default, the dataset returns a `PIL.Image` and a dictionary\n",
    "containing several fields, including `boxes`, `labels` and `masks`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RoAEkUgn4uEq"
   },
   "source": [
    "## Defining your model\n",
    "\n",
    "In this tutorial, we will be using [Mask R-CNN](https://arxiv.org/abs/1703.06870), which is based on top of [Faster R-CNN](https://arxiv.org/abs/1506.01497). Faster R-CNN is a model that predicts both bounding boxes and class scores for potential objects in the image.\n",
    "\n",
    "![Faster R-CNN](https://raw.githubusercontent.com/pytorch/vision/temp-tutorial/tutorials/tv_image03.png)\n",
    "\n",
    "Mask R-CNN adds an extra branch into Faster R-CNN, which also predicts segmentation masks for each instance.\n",
    "\n",
    "![Mask R-CNN](https://raw.githubusercontent.com/pytorch/vision/temp-tutorial/tutorials/tv_image04.png)\n",
    "\n",
    "There are two common situations where one might want to modify one of the available models in torchvision modelzoo.\n",
    "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).\n",
    "\n",
    "Let's go see how we would do one or another in the following sections.\n",
    "\n",
    "\n",
    "### 1 - Finetuning from a pretrained model\n",
    "\n",
    "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:\n",
    "```\n",
    "import torchvision\n",
    "from torchvision.models.detection.faster_rcnn import FastRCNNPredictor\n",
    "\n",
    "# load a model pre-trained pre-trained on COCO\n",
    "model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)\n",
    "\n",
    "# replace the classifier with a new one, that has\n",
    "# num_classes which is user-defined\n",
    "num_classes = 2  # 1 class (person) + background\n",
    "# get number of input features for the classifier\n",
    "in_features = model.roi_heads.box_predictor.cls_score.in_features\n",
    "# replace the pre-trained head with a new one\n",
    "model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n",
    "```\n",
    "\n",
    "### 2 - Modifying the model to add a different backbone\n",
    "\n",
    "Another common situation arises when the user wants to replace the backbone of a detection\n",
    "model with a different one. For example, the current default backbone (ResNet-50) might be too big for some applications, and smaller models might be necessary.\n",
    "\n",
    "Here is how we would go into leveraging the functions provided by torchvision to modify a backbone.\n",
    "\n",
    "```\n",
    "import torchvision\n",
    "from torchvision.models.detection import FasterRCNN\n",
    "from torchvision.models.detection.rpn import AnchorGenerator\n",
    "\n",
    "# load a pre-trained model for classification and return\n",
    "# only the features\n",
    "backbone = torchvision.models.mobilenet_v2(pretrained=True).features\n",
    "# FasterRCNN needs to know the number of\n",
    "# output channels in a backbone. For mobilenet_v2, it's 1280\n",
    "# so we need to add it here\n",
    "backbone.out_channels = 1280\n",
    "\n",
    "# let's make the RPN generate 5 x 3 anchors per spatial\n",
    "# location, with 5 different sizes and 3 different aspect\n",
    "# ratios. We have a Tuple[Tuple[int]] because each feature\n",
    "# map could potentially have different sizes and\n",
    "# aspect ratios\n",
    "anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),),\n",
    "                                   aspect_ratios=((0.5, 1.0, 2.0),))\n",
    "\n",
    "# let's define what are the feature maps that we will\n",
    "# use to perform the region of interest cropping, as well as\n",
    "# the size of the crop after rescaling.\n",
    "# if your backbone returns a Tensor, featmap_names is expected to\n",
    "# be [0]. More generally, the backbone should return an\n",
    "# OrderedDict[Tensor], and in featmap_names you can choose which\n",
    "# feature maps to use.\n",
    "roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],\n",
    "                                                output_size=7,\n",
    "                                                sampling_ratio=2)\n",
    "\n",
    "# put the pieces together inside a FasterRCNN model\n",
    "model = FasterRCNN(backbone,\n",
    "                   num_classes=2,\n",
    "                   rpn_anchor_generator=anchor_generator,\n",
    "                   box_roi_pool=roi_pooler)\n",
    "```\n",
    "\n",
    "### An Instance segmentation model for PennFudan Dataset\n",
    "\n",
    "In our case, we want to fine-tune from a pre-trained model, given that our dataset is very small. So we will be following approach number 1.\n",
    "\n",
    "Here we want to also compute the instance segmentation masks, so we will be using Mask R-CNN:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YjNHjVMOyYlH"
   },
   "outputs": [],
   "source": [
    "import torchvision\n",
    "from torchvision.models.detection.faster_rcnn import FastRCNNPredictor\n",
    "from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor\n",
    "\n",
    "\n",
    "def get_instance_segmentation_model(num_classes):\n",
    "    # load an instance segmentation model pre-trained on COCO\n",
    "    model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)\n",
    "\n",
    "    # get the number of input features for the classifier\n",
    "    in_features = model.roi_heads.box_predictor.cls_score.in_features\n",
    "    # replace the pre-trained head with a new one\n",
    "    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n",
    "\n",
    "    # now get the number of input features for the mask classifier\n",
    "    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels\n",
    "    hidden_layer = 256\n",
    "    # and replace the mask predictor with a new one\n",
    "    model.roi_heads.mask_predictor = MaskRCNNPredictor(\n",
    "        in_features_mask, hidden_layer, num_classes\n",
    "    )\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-WXLwePV5ieP"
   },
   "source": [
    "That's it, this will make model be ready to be trained and evaluated on our custom dataset.\n",
    "\n",
    "## Training and evaluation functions\n",
    "\n",
    "In `references/detection/,` we have a number of helper functions to simplify training and evaluating detection models.\n",
    "Here, we will use `references/detection/engine.py`, `references/detection/utils.py` and `references/detection/transforms.py`.\n",
    "\n",
    "Let's copy those files (and their dependencies) in here so that they are available in the notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "UYDb7PBw55b-",
    "outputId": "fc61f74a-e0c2-4a61-a25a-7c393871aaa8"
   },
   "outputs": [],
   "source": [
    "%%shell\n",
    "\n",
    "# Download TorchVision repo to use some files from\n",
    "# references/detection\n",
    "git clone https://github.com/pytorch/vision.git\n",
    "cd vision\n",
    "git checkout v0.8.2\n",
    "\n",
    "cp references/detection/utils.py ../\n",
    "cp references/detection/transforms.py ../\n",
    "cp references/detection/coco_eval.py ../\n",
    "cp references/detection/engine.py ../\n",
    "cp references/detection/coco_utils.py ../"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2u9e_pdv54nG"
   },
   "source": [
    "\n",
    "\n",
    "Let's write some helper functions for data augmentation / transformation, which leverages the functions in `refereces/detection` that we have just copied:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "l79ivkwKy357"
   },
   "outputs": [],
   "source": [
    "from engine import train_one_epoch, evaluate\n",
    "import utils\n",
    "import transforms as T\n",
    "\n",
    "\n",
    "def get_transform(train):\n",
    "    transforms = []\n",
    "    # converts the image, a PIL image, into a PyTorch Tensor\n",
    "    transforms.append(T.ToTensor())\n",
    "    if train:\n",
    "        # during training, randomly flip the training images\n",
    "        # and ground-truth for data augmentation\n",
    "        transforms.append(T.RandomHorizontalFlip(0.5))\n",
    "    return T.Compose(transforms)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ovjC81FZfWat"
   },
   "source": [
    "#### Testing forward() method\n",
    "\n",
    "Before iterating over the dataset, it’s good to see what the model expects during training and inference time on sample data.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "6-NFR--2fXV3",
    "outputId": "9b23511c-197d-4a36-a5a6-32d5bcb1faa1"
   },
   "outputs": [],
   "source": [
    "model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)\n",
    "dataset = PennFudanDataset(\"PennFudanPed\", get_transform(train=True))\n",
    "data_loader = torch.utils.data.DataLoader(\n",
    "    dataset, batch_size=2, shuffle=True, num_workers=4, collate_fn=utils.collate_fn\n",
    ")\n",
    "# For Training\n",
    "images, targets = next(iter(data_loader))\n",
    "images = list(image for image in images)\n",
    "targets = [{k: v for k, v in t.items()} for t in targets]\n",
    "output = model(images, targets)  # Returns losses and detections\n",
    "# For inference\n",
    "model.eval()\n",
    "x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]\n",
    "predictions = model(x)  # Returns predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FzCLqiZk-sjf"
   },
   "source": [
    "#### Note that we do not need to add a mean/std normalization nor image rescaling in the data transforms, as those are handled internally by the Mask R-CNN model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3YFJGJxk6XEs"
   },
   "source": [
    "### Putting everything together\n",
    "\n",
    "We now have the dataset class, the models and the data transforms. Let's instantiate them"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "a5dGaIezze3y",
    "outputId": "5cf908c4-a3c6-4c7b-b6f5-7633e8b43594"
   },
   "outputs": [],
   "source": [
    "# use our dataset and defined transformations\n",
    "dataset = PennFudanDataset(\"PennFudanPed\", get_transform(train=True))\n",
    "dataset_test = PennFudanDataset(\"PennFudanPed\", get_transform(train=False))\n",
    "\n",
    "# split the dataset in train and test set\n",
    "torch.manual_seed(1)\n",
    "indices = torch.randperm(len(dataset)).tolist()\n",
    "dataset = torch.utils.data.Subset(dataset, indices[:-50])\n",
    "dataset_test = torch.utils.data.Subset(dataset_test, indices[-50:])\n",
    "\n",
    "# define training and validation data loaders\n",
    "data_loader = torch.utils.data.DataLoader(\n",
    "    dataset, batch_size=2, shuffle=True, num_workers=4, collate_fn=utils.collate_fn\n",
    ")\n",
    "\n",
    "data_loader_test = torch.utils.data.DataLoader(\n",
    "    dataset_test,\n",
    "    batch_size=1,\n",
    "    shuffle=False,\n",
    "    num_workers=4,\n",
    "    collate_fn=utils.collate_fn,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "L5yvZUprj4ZN"
   },
   "source": [
    "Now let's instantiate the model and the optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zoenkCj18C4h"
   },
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "\n",
    "# our dataset has two classes only - background and person\n",
    "num_classes = 2\n",
    "\n",
    "# get the model using our helper function\n",
    "model = get_instance_segmentation_model(num_classes)\n",
    "# move model to the right device\n",
    "model.to(device)\n",
    "\n",
    "# construct an optimizer\n",
    "params = [p for p in model.parameters() if p.requires_grad]\n",
    "optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)\n",
    "\n",
    "# and a learning rate scheduler which decreases the learning rate by\n",
    "# 10x every 3 epochs\n",
    "lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XAd56lt4kDxc"
   },
   "source": [
    "And now let's train the model for 10 epochs, evaluating at the end of every epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "at-h4OWK0aoc",
    "outputId": "ab6753bc-6251-42e5-d65b-32efec93c61e"
   },
   "outputs": [],
   "source": [
    "# let's train it for 10 epochs\n",
    "from torch.optim.lr_scheduler import StepLR\n",
    "\n",
    "num_epochs = 10\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    # train for one epoch, printing every 10 iterations\n",
    "    train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10)\n",
    "    # update the learning rate\n",
    "    lr_scheduler.step()\n",
    "    # evaluate on the test dataset\n",
    "    evaluate(model, data_loader_test, device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Z6mYGFLxkO8F"
   },
   "source": [
    "Now that training has finished, let's have a look at what it actually predicts in a test image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YHwIdxH76uPj"
   },
   "outputs": [],
   "source": [
    "# pick one image from the test set\n",
    "img, _ = dataset_test[0]\n",
    "# put the model in evaluation mode\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    prediction = model([img.to(device)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DmN602iKsuey"
   },
   "source": [
    "Printing the prediction shows that we have a list of dictionaries. Each element of the list corresponds to a different image. As we have a single image, there is a single dictionary in the list.\n",
    "The dictionary contains the predictions for the image we passed. In this case, we can see that it contains `boxes`, `labels`, `masks` and `scores` as fields."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Lkmb3qUu6zw3",
    "outputId": "7e876402-3a0f-4696-a13e-6ddcabb5fb77"
   },
   "outputs": [],
   "source": [
    "prediction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RwT21rzotFbH"
   },
   "source": [
    "Let's inspect the image and the predicted segmentation masks.\n",
    "\n",
    "For that, we need to convert the image, which has been rescaled to 0-1 and had the channels flipped so that we have it in `[C, H, W]` format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 366
    },
    "id": "bpqN9t1u7B2J",
    "outputId": "b998eca7-5e3c-4ea6-ccfc-052fef26c5a0"
   },
   "outputs": [],
   "source": [
    "Image.fromarray(img.mul(255).permute(1, 2, 0).byte().numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "M58J3O9OtT1G"
   },
   "source": [
    "And let's now visualize the top predicted segmentation mask. The masks are predicted as `[N, 1, H, W]`, where `N` is the number of predictions, and are probability maps between 0-1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 366
    },
    "id": "5v5S3bm07SO1",
    "outputId": "89e92ecd-dbca-4939-fb94-8d2f1b3a7cc1"
   },
   "outputs": [],
   "source": [
    "Image.fromarray(prediction[0][\"masks\"][0, 0].mul(255).byte().cpu().numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0EZCVtCPunrT"
   },
   "source": [
    "Looks pretty good!\n",
    "\n",
    "## Wrapping up\n",
    "\n",
    "In this tutorial, you have learned how to create your own training pipeline for instance segmentation models, on a custom dataset.\n",
    "For that, you wrote a `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.\n",
    "\n",
    "For a more complete example, which includes multi-machine / multi-gpu training, check `references/detection/train.py`, which is present in the [torchvision GitHub repo](https://github.com/pytorch/vision/tree/v0.8.2/references/detection).\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "name": "torchvision_finetuning_instance_segmentation.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}