Object_Detection_Tutorial_(data_science_track).ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QKlO-kGhKk70"
      },
      "source": [
        "### Training a object detection Neural Network"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "First, upload data to Google Drive. Then, mount your Google Drive so that CoLab can have access to files in your Google Drive."
      ],
      "metadata": {
        "id": "Te1cQfV0-_CF"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mBsEUCJHCbSW",
        "outputId": "e2adfeed-7682-4600-a13f-21e76ea72f31"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Copy your dataset folder from the location in your Google Drive to current working directory in CoLab."
      ],
      "metadata": {
        "id": "QffHtAkw_N68"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "h0A4kuGUbhHZ"
      },
      "outputs": [],
      "source": [
        "!cp -R \"<<UPDATE THIS>>\" ."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t530KMEmNW4j"
      },
      "source": [
        "Install detecto. The reason we are using detecto is that its source code is very readable. It hides the nitty gritty details of typical training workflow in PyTorch."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mKRgxEtBQrt_"
      },
      "outputs": [],
      "source": [
        "!pip3 install detecto"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "zhU64wlfDi-v"
      },
      "outputs": [],
      "source": [
        "from detecto import core, utils, config\n",
        "import torch\n",
        "import torchvision\n",
        "from torchvision.models.detection.faster_rcnn import FastRCNNPredictor\n",
        "from torchvision import transforms\n",
        "import matplotlib.pyplot as plt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "r7Wabhmus0LZ"
      },
      "outputs": [],
      "source": [
        "class FreezableModel(core.Model):\n",
        "  \"\"\"\n",
        "    Subclass Detecto core.Model to allow user to specify which layers of back-\n",
        "    bone they want to freeze \n",
        "  \"\"\"\n",
        "  DEFAULT = 'fasterrcnn_resnet50_fpn'\n",
        "  MOBILENET = 'fasterrcnn_mobilenet_v3_large_fpn'\n",
        "  MOBILENET_320 = 'fasterrcnn_mobilenet_v3_large_320_fpn'\n",
        "\n",
        "  def __init__(self, classes=None, device=None, pretrained=True,\n",
        "                model_name=DEFAULT, trainable_backbone_layers=None):\n",
        "    self._device = device if device else config.config['default_device']\n",
        "\n",
        "    # Load a model pre-trained on COCO\n",
        "    if model_name == self.DEFAULT:\n",
        "        self._model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=pretrained,\n",
        "                                                                            trainable_backbone_layers=trainable_backbone_layers)\n",
        "    elif model_name == self.MOBILENET:\n",
        "        self._model = torchvision.models.detection.fasterrcnn_mobilenet_v3_large_fpn(pretrained=pretrained,\n",
        "                                                                                      trainable_backbone_layers=trainable_backbone_layers)\n",
        "    elif model_name == self.MOBILENET_320:\n",
        "        self._model = torchvision.models.detection.fasterrcnn_mobilenet_v3_large_320_fpn(pretrained=pretrained,\n",
        "                                                                                          trainable_backbone_layers=trainable_backbone_layers)\n",
        "    else:\n",
        "        raise ValueError(f'Invalid value {model_name} for model_name. ' +\n",
        "                          f'Please choose between {self.DEFAULT}, {self.MOBILENET}, and {self.MOBILENET_320}.')\n",
        "\n",
        "    if classes:\n",
        "        # Get the number of input features for the classifier\n",
        "        in_features = self._model.roi_heads.box_predictor.cls_score.in_features\n",
        "        # Replace the pre-trained head with a new one (note: +1 because of the __background__ class)\n",
        "        self._model.roi_heads.box_predictor = FastRCNNPredictor(in_features, len(classes) + 1)\n",
        "        self._disable_normalize = False\n",
        "    else:\n",
        "        classes = config['default_classes']\n",
        "        self._disable_normalize = True\n",
        "\n",
        "    self._model.to(self._device)\n",
        "\n",
        "    # Mappings to convert from string labels to ints and vice versa\n",
        "    self._classes = ['__background__'] + classes\n",
        "    self._int_mapping = {label: index for index, label in enumerate(self._classes)}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cUNZl-v1Djxa"
      },
      "outputs": [],
      "source": [
        "# Convert XML files to CSV format\n",
        "# You may have to update this based on how you have placed your train and val labels\n",
        "utils.xml_to_csv('dataset/train_labels/', 'train_labels.csv') \n",
        "utils.xml_to_csv('dataset/val_labels/', 'val_labels.csv')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {
        "id": "q3cDsGRPDzc3"
      },
      "outputs": [],
      "source": [
        "# Define custom transforms to apply to your dataset\n",
        "custom_transforms = transforms.Compose([\n",
        "    transforms.ToPILImage(),\n",
        "    transforms.Resize(500),\n",
        "    transforms.ColorJitter(brightness=[0.9,1.1], contrast=[0.9,1.1], saturation=0.3),\n",
        "    transforms.GaussianBlur(kernel_size=(5, 5), sigma=(0.1, 0.3)),\n",
        "    transforms.RandomPosterize(bits=6),\n",
        "    transforms.RandomAdjustSharpness(sharpness_factor=2),\n",
        "    transforms.RandomAutocontrast(),\n",
        "    transforms.RandomEqualize(),\n",
        "    transforms.ToTensor(),\n",
        "    utils.normalize_transform(),\n",
        "])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {
        "id": "E8_IkZ55D4Zz"
      },
      "outputs": [],
      "source": [
        "# Pass in a CSV file instead of XML files for faster Dataset initialization speeds\n",
        "# You may have to update dataset/train_images and dataset/val_images based on \n",
        "# how you have placed your images\n",
        "dataset = core.Dataset('train_labels.csv', 'dataset/train_images/', transform=custom_transforms)\n",
        "val_dataset = core.Dataset('val_labels.csv', 'dataset/val_images/')  # Validation dataset for training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "metadata": {
        "id": "btGAl41pD9_W"
      },
      "outputs": [],
      "source": [
        "# Create your own DataLoader with custom options\n",
        "loader = core.DataLoader(dataset, batch_size=4, shuffle=True) "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wPTfrrENcVTb"
      },
      "source": [
        "Visulize images and labels from dataloader"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "metadata": {
        "id": "SPGsioiJQ_Zp"
      },
      "outputs": [],
      "source": [
        "import cv2\n",
        "import numpy as np"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "K2_4tGO1Ov21"
      },
      "outputs": [],
      "source": [
        "images, labels = next(iter(loader))\n",
        "for img, label in zip(images, labels):\n",
        "  bboxes = label['boxes']\n",
        "  names = label['labels']\n",
        "  # img = utils.reverse_normalize(img)\n",
        "  img = np.ascontiguousarray(img.permute(1,2,0))\n",
        "  img = np.uint8(img*255)\n",
        "  for bbox, name in zip(bboxes, names):\n",
        "    print(bbox, name)\n",
        "    bbox = bbox.numpy()\n",
        "    start_pt = (bbox[0], bbox[1])\n",
        "    end_pt = (bbox[2], bbox[3])\n",
        "    img = cv2.rectangle(img, start_pt, end_pt, (0,255,255), 5)\n",
        "    img = cv2.putText(img, name, (bbox[0]-10, bbox[1]-10), \n",
        "                      cv2.FONT_HERSHEY_SIMPLEX, 2, (255,0,255), 5, cv2.LINE_AA)\n",
        "\n",
        "  plt.imshow(img)\n",
        "  plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "metadata": {
        "id": "wb64FmHXEA9s"
      },
      "outputs": [],
      "source": [
        "class_names = np.unique(dataset._csv['class'])\n",
        "print(f\"Classes in dataset: {class_names}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 38,
      "metadata": {
        "id": "yJQ1Jmd9RXqd"
      },
      "outputs": [],
      "source": [
        "model = FreezableModel(list(class_names), model_name='fasterrcnn_resnet50_fpn', trainable_backbone_layers=0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dOGuP7qfUM1t"
      },
      "outputs": [],
      "source": [
        "for name, parameter in model._model.named_parameters():\n",
        "  print(name, parameter.requires_grad)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9WyDrJmWRU35"
      },
      "outputs": [],
      "source": [
        "losses = model.fit(loader, val_dataset, epochs=10, learning_rate=0.008, verbose=True)\n",
        "\n",
        "plt.plot(losses)  # Visualize loss throughout training\n",
        "plt.show()\n",
        "\n",
        "model.save(\"<<UPDATE THIS>>/model_weights.pth\")  # Save model to a file\n",
        "# Update this path so that the weights are saved on your Google Drive"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zA1uk54yFllt"
      },
      "outputs": [],
      "source": [
        "# Directly access underlying torchvision model for even more control\n",
        "torch_model = model.get_internal_model()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zg0ww75pKNda"
      },
      "source": [
        "## Take photo from web cam and perform inference with trained model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 40,
      "metadata": {
        "id": "SvBTnqVdIw78"
      },
      "outputs": [],
      "source": [
        "from IPython.display import display, Javascript\n",
        "from google.colab.output import eval_js\n",
        "from base64 import b64decode\n",
        "\n",
        "def take_photo(filename='photo.jpg', quality=0.8):\n",
        "  js = Javascript('''\n",
        "    async function takePhoto(quality) {\n",
        "      const div = document.createElement('div');\n",
        "      const capture = document.createElement('button');\n",
        "      capture.textContent = 'Capture';\n",
        "      div.appendChild(capture);\n",
        "\n",
        "      const video = document.createElement('video');\n",
        "      video.style.display = 'block';\n",
        "      const stream = await navigator.mediaDevices.getUserMedia({video: true});\n",
        "\n",
        "      document.body.appendChild(div);\n",
        "      div.appendChild(video);\n",
        "      video.srcObject = stream;\n",
        "      await video.play();\n",
        "\n",
        "      // Resize the output to fit the video element.\n",
        "      google.colab.output.setIframeHeight(document.documentElement.scrollHeight, true);\n",
        "\n",
        "      // Wait for Capture to be clicked.\n",
        "      await new Promise((resolve) => capture.onclick = resolve);\n",
        "\n",
        "      const canvas = document.createElement('canvas');\n",
        "      canvas.width = video.videoWidth;\n",
        "      canvas.height = video.videoHeight;\n",
        "      canvas.getContext('2d').drawImage(video, 0, 0);\n",
        "      stream.getVideoTracks()[0].stop();\n",
        "      div.remove();\n",
        "      return canvas.toDataURL('image/jpeg', quality);\n",
        "    }\n",
        "    ''')\n",
        "  display(js)\n",
        "  data = eval_js('takePhoto({})'.format(quality))\n",
        "  binary = b64decode(data.split(',')[1])\n",
        "  with open(filename, 'wb') as f:\n",
        "    f.write(binary)\n",
        "  return filename"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fQv8K3kII4Qq"
      },
      "outputs": [],
      "source": [
        "from IPython.display import Image\n",
        "try:\n",
        "  filename = take_photo()\n",
        "  print('Saved to {}'.format(filename))\n",
        "  \n",
        "  # Show the image which was just taken.\n",
        "  display(Image(filename))\n",
        "except Exception as err:\n",
        "  # Errors will be thrown if the user does not have a webcam or if they do not\n",
        "  # grant the page permission to access it.\n",
        "  print(str(err))"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Perform inference"
      ],
      "metadata": {
        "id": "wAFrhEM5AC0z"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 42,
      "metadata": {
        "id": "WV_rZWsQn5IF"
      },
      "outputs": [],
      "source": [
        "from detecto import visualize"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Jy7LEZ9zKWXQ"
      },
      "outputs": [],
      "source": [
        "model = FreezableModel.load(\"<<UPDATE THIS>>/model_weights.pth\", list(class_names))\n",
        "image = utils.read_image('/content/photo.jpg')\n",
        "model._disable_normalize = True \n",
        "labels, boxes, scores = model.predict(image)  # Get all predictions on an image\n",
        "\n",
        "# Get top prediction\n",
        "idx_score = np.argmax(scores.numpy())\n",
        "box = boxes[idx_score]\n",
        "label = labels[idx_score]\n",
        "\n",
        "visualize.show_labeled_image(image, box, label)  # Plot top prediction"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "REqWtsDaRQ-U"
      },
      "source": [
        "## Fine tuning some later stages of backbone"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 46,
      "metadata": {
        "id": "NnV107QEGHlL"
      },
      "outputs": [],
      "source": [
        "trainable_layers = 3"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 47,
      "metadata": {
        "id": "1ehEjvGwFlWI"
      },
      "outputs": [],
      "source": [
        "if trainable_layers < 0 or trainable_layers > 5:\n",
        "    raise ValueError(f\"Trainable layers should be in the range [0,5], got {trainable_layers}\")\n",
        "layers_to_train = [\"fpn\", \"body.layer4\", \"body.layer3\", \"body.layer2\", \"body.layer1\", \"body.conv1\"][:trainable_layers]\n",
        "if trainable_layers == 5:\n",
        "    layers_to_train.append(\"bn1\")\n",
        "for name, parameter in model._model.backbone.named_parameters():\n",
        "    if all([not name.startswith(layer) for layer in layers_to_train]):\n",
        "        parameter.requires_grad = False"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MrvVtoVLLWhW"
      },
      "outputs": [],
      "source": [
        "for name, parameter in model._model.backbone.named_parameters():\n",
        "  print(name, parameter.requires_grad)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GtUIy01HOXjM"
      },
      "outputs": [],
      "source": [
        "losses = model.fit(loader, val_dataset, epochs=5, learning_rate=0.003, verbose=True)\n",
        "\n",
        "plt.plot(losses)  # Visualize loss throughout training\n",
        "plt.show()\n",
        "\n",
        "model.save(\"<<UPDATE THIS>>/fine_tuned_model_weights.pth\")  # Save fine-tuned model to a file"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "Object Detection Tutorial (data science track).ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}