Demo_Pytorch_MaskRCNN_With_Video.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cWN0o-mcXmvs"
      },
      "source": [
        "# **Demo Mask R-CNN with COCO dataset**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ooKMPGYxXqvu"
      },
      "source": [
        "## Connect to drive"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FwbLEM0IXsqu",
        "outputId": "aca755a4-aaab-461d-df0f-cf1a02e3d528"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/\n",
            "Mounted at content/gdrive\n"
          ]
        }
      ],
      "source": [
        "%cd ..\n",
        "from google.colab import drive\n",
        "drive.mount('content/gdrive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qh1XXliIXtQ_"
      },
      "outputs": [],
      "source": [
        "!ln -s /content/gdrive/My\\ Drive/ /colab"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "M113aSZ_XupX",
        "outputId": "63778110-82dc-479d-9848-8ae8ff0df3f8"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/content/gdrive/My Drive/TrainPBL-customData\n"
          ]
        }
      ],
      "source": [
        "cd /colab/TrainPBL-customData"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wpQ9tmiSX1WM"
      },
      "source": [
        "## Set up evironment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xYJHvYTJX4vm"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "# Install PyTorch with CUDA 11.8\n",
        "# %pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 torchtext==0.14.1 torchaudio==0.13.1 torchdata==0.5.1 --extra-index-url https://download.pytorch.org/whl/cu117\n",
        "!pip install torch==1.13.1+cu118 torchvision==0.14.1+cu118 torchaudio==0.13.1+cu118 --index-url https://download.pytorch.org/whl/cu118\n",
        "\n",
        "# Install additional dependencies\n",
        "!pip install matplotlib pandas pillow torchtnt==0.2.0 tqdm tabulate\n",
        "\n",
        "# Install package for creating visually distinct colormaps\n",
        "!pip install distinctipy\n",
        "\n",
        "# Install utility packages\n",
        "!pip install cjm_pandas_utils cjm_psl_utils cjm_pil_utils cjm_pytorch_utils cjm_torchvision_tfms"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Y2YFphNPX6Mt",
        "outputId": "48fc249c-4389-4dfd-9f23-e369a6db4aee"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: pycocotools in /usr/local/lib/python3.10/dist-packages (2.0.7)\n",
            "Requirement already satisfied: matplotlib>=2.1.0 in /usr/local/lib/python3.10/dist-packages (from pycocotools) (3.7.1)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from pycocotools) (1.25.2)\n",
            "Requirement already satisfied: contourpy>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=2.1.0->pycocotools) (1.2.1)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=2.1.0->pycocotools) (0.12.1)\n",
            "Requirement already satisfied: fonttools>=4.22.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=2.1.0->pycocotools) (4.53.0)\n",
            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=2.1.0->pycocotools) (1.4.5)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=2.1.0->pycocotools) (24.0)\n",
            "Requirement already satisfied: pillow>=6.2.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=2.1.0->pycocotools) (9.4.0)\n",
            "Requirement already satisfied: pyparsing>=2.3.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=2.1.0->pycocotools) (3.1.2)\n",
            "Requirement already satisfied: python-dateutil>=2.7 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=2.1.0->pycocotools) (2.8.2)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil>=2.7->matplotlib>=2.1.0->pycocotools) (1.16.0)\n",
            "Requirement already satisfied: opencv-python in /usr/local/lib/python3.10/dist-packages (4.8.0.76)\n",
            "Requirement already satisfied: numpy>=1.21.2 in /usr/local/lib/python3.10/dist-packages (from opencv-python) (1.25.2)\n"
          ]
        }
      ],
      "source": [
        "!pip install pycocotools\n",
        "!pip install opencv-python"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6xLMBnrJX8lk"
      },
      "source": [
        "## Importing the Required Dependencies"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ulkfRagFX_TU"
      },
      "outputs": [],
      "source": [
        "# Import Python Standard Library dependencies\n",
        "import datetime\n",
        "from functools import partial\n",
        "from glob import glob\n",
        "import json\n",
        "import math\n",
        "import multiprocessing\n",
        "import os\n",
        "from pathlib import Path\n",
        "import random\n",
        "from typing import Any, Dict, Optional\n",
        "\n",
        "# Import utility functions\n",
        "from cjm_psl_utils.core import download_file, file_extract, get_source_code\n",
        "from cjm_pil_utils.core import resize_img, get_img_files, stack_imgs\n",
        "from cjm_pytorch_utils.core import pil_to_tensor, tensor_to_pil, get_torch_device, set_seed, denorm_img_tensor, move_data_to_device\n",
        "from cjm_pandas_utils.core import markdown_to_pandas, convert_to_numeric, convert_to_string\n",
        "from cjm_torchvision_tfms.core import ResizeMax, PadSquare, CustomRandomIoUCrop\n",
        "\n",
        "# Import the distinctipy module\n",
        "from distinctipy import distinctipy\n",
        "\n",
        "# Import matplotlib for creating plots\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "# Import numpy\n",
        "import numpy as np\n",
        "\n",
        "# Import the pandas package\n",
        "import pandas as pd\n",
        "\n",
        "# Set options for Pandas DataFrame display\n",
        "pd.set_option('max_colwidth', None)  # Do not truncate the contents of cells in the DataFrame\n",
        "pd.set_option('display.max_rows', None)  # Display all rows in the DataFrame\n",
        "pd.set_option('display.max_columns', None)  # Display all columns in the DataFrame\n",
        "\n",
        "# Import PIL for image manipulation\n",
        "from PIL import Image, ImageDraw\n",
        "\n",
        "# Import PyTorch dependencies\n",
        "import torch\n",
        "from torch.amp import autocast\n",
        "from torch.cuda.amp import GradScaler\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "from torchtnt.utils import get_module_summary\n",
        "import torchvision\n",
        "torchvision.disable_beta_transforms_warning()\n",
        "from torchvision.tv_tensors import BoundingBoxes, Mask\n",
        "from torchvision.utils import draw_bounding_boxes, draw_segmentation_masks\n",
        "import torchvision.transforms.v2  as transforms\n",
        "from torchvision.transforms.v2 import functional as TF\n",
        "\n",
        "# Import Mask R-CNN\n",
        "from torchvision.models.detection import maskrcnn_resnet50_fpn_v2, MaskRCNN\n",
        "from torchvision.models.detection import MaskRCNN_ResNet50_FPN_V2_Weights\n",
        "from torchvision.models.detection.faster_rcnn import FastRCNNPredictor\n",
        "from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor\n",
        "\n",
        "# Import tqdm for progress bar\n",
        "from tqdm.auto import tqdm\n",
        "from pycocotools import mask as coco_mask\n",
        "from pycocotools.coco import COCO\n",
        "from pycocotools.cocoeval import COCOeval\n",
        "from PIL import Image, ImageDraw\n",
        "from torchvision.transforms.v2.functional import to_tensor, to_pil_image\n",
        "\n",
        "import cv2\n",
        "import os\n",
        "import time"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6ogeCGeLYE0Q",
        "outputId": "c5beb7b8-ba33-4271-819c-77df8b471646"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "('cuda', torch.float32)"
            ]
          },
          "metadata": {},
          "execution_count": 7
        }
      ],
      "source": [
        "# Set the seed for generating random numbers in PyTorch, NumPy, and Python's random module.\n",
        "seed = 1234\n",
        "set_seed(seed)\n",
        "device = get_torch_device()\n",
        "dtype = torch.float32\n",
        "device, dtype"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lER4GDWuYLfQ"
      },
      "source": [
        "## Setting the Model Paths"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "B8c0Zanz_nBP",
        "outputId": "bddbac10-ca42-4c04-e9cd-028ae724ee6f"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/content/gdrive/MyDrive/TrainPBL-customData\n"
          ]
        }
      ],
      "source": [
        "!pwd"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vWoGGvDVYKyi"
      },
      "outputs": [],
      "source": [
        "# /content/gdrive/MyDrive/TrainPBL-customData/pytorch-mask-r-cnn-instance-segmentation-custom-data/2024-06-09_05-23-18/maskrcnn_resnet50_fpn_v2.pth\n",
        "model_path = './pytorch-mask-r-cnn-instance-segmentation-custom-data/2024-06-09_05-23-18/maskrcnn_resnet50_fpn_v2.pth'\n",
        "video_path = './coco/video_test/input-video-testing.mp4'\n",
        "output_video_path = './coco/video_test/output_video.avi'\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Jkv197Et-mxb"
      },
      "source": [
        "## Set up Class name"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OejG-ElFdtRo",
        "outputId": "1f499b9e-0f9a-4c7e-c2a2-76ae00d59cde"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "       Class Names\n",
            "0       background\n",
            "1              dog\n",
            "2     potted plant\n",
            "3               tv\n",
            "4             bird\n",
            "5              cat\n",
            "6            horse\n",
            "7            sheep\n",
            "8              cow\n",
            "9           bottle\n",
            "10           couch\n",
            "11           chair\n",
            "12    dining table\n",
            "13         bicycle\n",
            "14             car\n",
            "15      motorcycle\n",
            "16        airplane\n",
            "17             bus\n",
            "18           train\n",
            "19            boat\n",
            "20          person\n",
            "21       stop sign\n",
            "22        umbrella\n",
            "23             tie\n",
            "24        sandwich\n",
            "25             bed\n",
            "26      cell phone\n",
            "27    refrigerator\n",
            "28           clock\n",
            "29      toothbrush\n",
            "30           truck\n",
            "31   traffic light\n",
            "32    fire hydrant\n",
            "33   parking meter\n",
            "34           bench\n",
            "35         frisbee\n",
            "36            skis\n",
            "37       snowboard\n",
            "38      skateboard\n",
            "39       surfboard\n",
            "40      wine glass\n",
            "41             cup\n",
            "42            fork\n",
            "43           knife\n",
            "44           spoon\n",
            "45            bowl\n",
            "46          banana\n",
            "47           apple\n",
            "48          orange\n",
            "49        broccoli\n",
            "50          carrot\n",
            "51           pizza\n",
            "52            cake\n",
            "53          toilet\n",
            "54          laptop\n",
            "55           mouse\n",
            "56          remote\n",
            "57        keyboard\n",
            "58       microwave\n",
            "59            oven\n",
            "60         toaster\n",
            "61            sink\n",
            "62            book\n",
            "63            vase\n",
            "64        scissors\n",
            "65      teddy bear\n",
            "66      hair drier\n",
            "67        backpack\n",
            "68         handbag\n",
            "69        suitcase\n",
            "70        elephant\n",
            "71           zebra\n",
            "72            kite\n",
            "73         hot dog\n",
            "74           donut\n",
            "75         giraffe\n",
            "76    baseball bat\n",
            "77  baseball glove\n",
            "78     sports ball\n",
            "79            bear\n",
            "80   tennis racket\n"
          ]
        }
      ],
      "source": [
        "# Danh sách các nhãn lớp\n",
        "labels = [\n",
        "    'background', 'dog', 'potted plant', 'tv', 'bird', 'cat', 'horse', 'sheep', 'cow', 'bottle', 'couch', 'chair', 'dining table',\n",
        "    'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'boat', 'person', 'stop sign', 'umbrella', 'tie', 'sandwich',\n",
        "    'bed', 'cell phone', 'refrigerator', 'clock', 'toothbrush', 'truck', 'traffic light', 'fire hydrant', 'parking meter',\n",
        "    'bench', 'frisbee', 'skis', 'snowboard', 'skateboard', 'surfboard', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',\n",
        "    'banana', 'apple', 'orange', 'broccoli', 'carrot', 'pizza', 'cake', 'toilet', 'laptop', 'mouse', 'remote', 'keyboard',\n",
        "    'microwave', 'oven', 'toaster', 'sink', 'book', 'vase', 'scissors', 'teddy bear', 'hair drier', 'backpack', 'handbag',\n",
        "    'suitcase', 'elephant', 'zebra', 'kite', 'hot dog', 'donut', 'giraffe', 'baseball bat', 'baseball glove', 'sports ball',\n",
        "    'bear', 'tennis racket'\n",
        "]\n",
        "\n",
        "# Tạo DataFrame từ danh sách các nhãn lớp\n",
        "temp_df = pd.DataFrame(labels, columns=[\"Class Names\"])\n",
        "\n",
        "# Chuyển đổi DataFrame thành danh sách và thêm nhãn 'background'\n",
        "class_names = temp_df[\"Class Names\"].tolist()\n",
        "\n",
        "# Hiển thị các nhãn lớp để kiểm tra\n",
        "print(pd.DataFrame(class_names, columns=[\"Class Names\"]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jFjmW68YZCfZ"
      },
      "source": [
        "## Loading the Mask R-CNN Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gS7eTZRLZJvc",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "cb44ebf2-a221-4867-c9b0-52a9df689c77"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Downloading: \"https://download.pytorch.org/models/maskrcnn_resnet50_fpn_v2_coco-73cbd019.pth\" to /root/.cache/torch/hub/checkpoints/maskrcnn_resnet50_fpn_v2_coco-73cbd019.pth\n",
            "100%|██████████| 177M/177M [00:01<00:00, 101MB/s]\n"
          ]
        }
      ],
      "source": [
        "# Initialize the Mask R-CNN model with pretrained weights\n",
        "model = maskrcnn_resnet50_fpn_v2(weights='DEFAULT')\n",
        "\n",
        "# Get the number of input features for the classifier and mask predictor\n",
        "in_features_box = model.roi_heads.box_predictor.cls_score.in_features\n",
        "in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels\n",
        "\n",
        "# Get the number of output channels for the Mask Predictor\n",
        "dim_reduced = model.roi_heads.mask_predictor.conv5_mask.out_channels\n",
        "\n",
        "# Replace the box predictor\n",
        "model.roi_heads.box_predictor = FastRCNNPredictor(in_channels=in_features_box, num_classes=len(class_names))\n",
        "\n",
        "# Replace the mask predictor\n",
        "model.roi_heads.mask_predictor = MaskRCNNPredictor(in_channels=in_features_mask, dim_reduced=dim_reduced, num_classes=len(class_names))\n",
        "\n",
        "# Set the model's device and data type\n",
        "model.to(device=device, dtype=dtype)\n",
        "\n",
        "# Load the model state from the file\n",
        "model.load_state_dict(torch.load(model_path, map_location=device))\n",
        "\n",
        "# Add attributes to store the device and model name for later reference\n",
        "model.device = device\n",
        "model.name = 'maskrcnn_resnet50_fpn_v2'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 135
        },
        "id": "3wZ8xmB7Zz2A",
        "outputId": "d1851fd2-54e6-428e-e6e6-f3dc6f7d42b7"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/torchtnt/utils/module_summary.py:271: UserWarning: Backward FLOPs are only computed if module foward returns a tensor.\n",
            "  warnings.warn(\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "       Type # Parameters # Trainable Parameters Size (bytes) Forward FLOPs\n",
              "0  MaskRCNN       46.3 M                 46.1 M        185 M         285 G"
            ],
            "text/html": [
              "\n",
              "  <div id=\"df-6f4ec53c-4fc3-4ea0-acce-496667984982\" class=\"colab-df-container\">\n",
              "    <div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>Type</th>\n",
              "      <th># Parameters</th>\n",
              "      <th># Trainable Parameters</th>\n",
              "      <th>Size (bytes)</th>\n",
              "      <th>Forward FLOPs</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>MaskRCNN</td>\n",
              "      <td>46.3 M</td>\n",
              "      <td>46.1 M</td>\n",
              "      <td>185 M</td>\n",
              "      <td>285 G</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n",
              "    <div class=\"colab-df-buttons\">\n",
              "\n",
              "  <div class=\"colab-df-container\">\n",
              "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-6f4ec53c-4fc3-4ea0-acce-496667984982')\"\n",
              "            title=\"Convert this dataframe to an interactive table.\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
              "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
              "  </svg>\n",
              "    </button>\n",
              "\n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\n",
              "      gap: 12px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert {\n",
              "      background-color: #E8F0FE;\n",
              "      border: none;\n",
              "      border-radius: 50%;\n",
              "      cursor: pointer;\n",
              "      display: none;\n",
              "      fill: #1967D2;\n",
              "      height: 32px;\n",
              "      padding: 0 0 0 0;\n",
              "      width: 32px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert:hover {\n",
              "      background-color: #E2EBFA;\n",
              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "      fill: #174EA6;\n",
              "    }\n",
              "\n",
              "    .colab-df-buttons div {\n",
              "      margin-bottom: 4px;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert {\n",
              "      background-color: #3B4455;\n",
              "      fill: #D2E3FC;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert:hover {\n",
              "      background-color: #434B5C;\n",
              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "      fill: #FFFFFF;\n",
              "    }\n",
              "  </style>\n",
              "\n",
              "    <script>\n",
              "      const buttonEl =\n",
              "        document.querySelector('#df-6f4ec53c-4fc3-4ea0-acce-496667984982 button.colab-df-convert');\n",
              "      buttonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "      async function convertToInteractive(key) {\n",
              "        const element = document.querySelector('#df-6f4ec53c-4fc3-4ea0-acce-496667984982');\n",
              "        const dataTable =\n",
              "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
              "                                                    [key], {});\n",
              "        if (!dataTable) return;\n",
              "\n",
              "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
              "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
              "          + ' to learn more about interactive tables.';\n",
              "        element.innerHTML = '';\n",
              "        dataTable['output_type'] = 'display_data';\n",
              "        await google.colab.output.renderOutput(dataTable, element);\n",
              "        const docLink = document.createElement('div');\n",
              "        docLink.innerHTML = docLinkHtml;\n",
              "        element.appendChild(docLink);\n",
              "      }\n",
              "    </script>\n",
              "  </div>\n",
              "\n",
              "\n",
              "    </div>\n",
              "  </div>\n"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "dataframe",
              "summary": "{\n  \"name\": \"summary_df\",\n  \"rows\": 1,\n  \"fields\": [\n    {\n      \"column\": \"Type\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 1,\n        \"samples\": [\n          \"MaskRCNN\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"# Parameters\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 1,\n        \"samples\": [\n          \"46.3 M\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"# Trainable Parameters\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 1,\n        \"samples\": [\n          \"46.1 M\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"Size (bytes)\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 1,\n        \"samples\": [\n          \"185 M\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"Forward FLOPs\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 1,\n        \"samples\": [\n          \"285 G\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    }\n  ]\n}"
            }
          },
          "metadata": {},
          "execution_count": 12
        }
      ],
      "source": [
        "test_inp = torch.randn(1, 3, 256, 256).to(device)\n",
        "\n",
        "summary_df = markdown_to_pandas(f\"{get_module_summary(model.eval(), [test_inp])}\")\n",
        "\n",
        "# # Filter the summary to only contain Conv2d layers and the model\n",
        "summary_df = summary_df[summary_df.index == 0]\n",
        "\n",
        "summary_df.drop(['In size', 'Out size', 'Contains Uninitialized Parameters?'], axis=1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 653
        },
        "id": "82XpMk21edS0",
        "outputId": "497e356d-5760-44f4-a104-a9ac4d2f5f7d"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 800x800 with 1 Axes>"
            ],
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAj8AAAJ8CAYAAAARYMD1AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAX2klEQVR4nO3aa6ikeWLX8V/dzr1P98z0TG92dnazu4ZRk8VL2GBUAsYQBO+KBAwhCokazEp8ZSCCQXwjRBCJCCYx4IskKhgQ88p4SQxJwIAuuSy7kk3ovWV2err7nO5zqzpV5QsJeb3nDM8zze/zef9/+ne6nqr61jk12W632wAAlJiOPQAAYEjiBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoMr/tBR7mYR7l0bux5YXw4OGDvP7o9bFnDOrhXvLopbFXDOfB9cO8Pu+5p5PkZL6X872eB/nyd5I8X4w9YzCnu4+zeely7BmDund4kQd3eh7j1cNltu/sjD1jUDv37+fggx+80dlbxc/DPMybeTOX6XhSvfHwjXzmzc+k5MdNkjy8m7z5ieTy1pn8YngjD/OZvJmmB/lk925++Bs+ketpx4O8+vIsv/k335/tajL2lEGcHz/Nz37PD2Uzvx57ymAeHK/z77/ncU7WYy8ZxvXD5Asfm2V7NfaSYU339vItv/HrNwqgW/3Z61Ee1YRPktx/dD97l3tjzxjUo4Oe8EmS+3mUvaJ7OknOFwc14ZMk1yfTmvBJkuXBWVX4JMm9g012e27prN9JXfgkyebyMstHN/stve/8AABVxA8AUEX8AABVxA8AUEX8AABVxA8AUEX8AABVxA8AUEX8AABVxA8AUEX8AABVxA8AUEX8AABVxA8AUEX8AABVxA8AUEX8AABVxA8AUEX8AABVxA8AUEX8AABVxA8AUEX8AABVxA8AUEX8AABVxA8AUEX8AABVxA8AUEX8AABVxA8AUEX8AABVxA8AUEX8AABVxA8AUEX8AABVxA8AUEX8AABVxA8AUEX8AABVxA8AUEX8AABVxA8AUEX8AABVxA8AUEX8AABVxA8AUEX8AABVxA8AUEX8AABVxA8AUEX8AABVxA8AUEX8AABVxA8AUEX8AABVxA8AUEX8AABV5re9wP0nb+T47P67seU9732Pf38+/Qe2OdubjD1lMJ96NcnrJ8n6fOwpgzjP0zx89Hre92Q99pTBbK9fyfL/HufR9mDsKYO4ejjPO/u7Wew9H3vKIM6OkntXx9mfdzyHk+Tly+TpOx/O5Oju2FMGsZol8z94ncOL07GnDGr57OzGZ28VP8dPHuRf/tNPZ+d67zaXeWG8fS/52CeT1WLsJUM6SbY/nEyuxx4yiGdvbbP69r+Rt1djLxnOFy7v5p/9yvfmatNxY08OTnP8F/91JrOOwD2eLvN9l3cynxyOPWUw28mD/PKT/5A82x17yiDu7D3L3/rHP5n5tuOe/l3r5Sqr45t9aLtV/Nw/e392Ot4TkySnh8lq0fNbn//vvCZ8kuTlk2SvKHyS5PHqoCZ8kmS6d14TPklyML3KfLIZe8aw9u4l847wSZKD9WVd+CTJbGeR2f2Xb3TWd34AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoIn4AgCriBwCoMr/tBc73P5fl7qN3Y8t73sVB8vpvvS/T51819pQBzXP31aMs3vd87CGD+H27yZc/vJN3zm/91Hhh/Padnbz0+FH2zsdeMozp8dPsL8+yXF2OPWUQ+/OT7E1XOZ5ux54ymO30NJ9/61fz6N7h2FMGsVo9ySefT3J4eWfsKYM63lvmtRuevdUr/POdz+W/fusfymZ2dZvLvDB2Hn0gP/DN/zvT1WbsKYM5ev8s3/bzF7ev5BfE51/byZ/461+Xq03PL0WPLq7ziZ/7j5mX3Nari9P81i/+dLab9dhTBvFK1vmunUfZmYy9ZDifv/ck3/jN/yhXFx3P453Te/mFf/cPMl0vxp4yqMVsnR/5pkUe3PvKz97qzrg+eKcmfJJkfvpKpqu9sWcMau/l88z3Ot4kkuTxcl4VPkmyt9zUhE+SrFeXNeGTJEeTTVX4JMnjw0WuFj3P4/nFYV34JMlqPcuz1f6NzvbcHQAAET8AQBnxAwBUET8AQBXxAwBUET8AQBXxAwBUET8AQBXxAwBUET8AQBXxAwBUET8AQBXxAwBUET8AQBXxAwBUET8AQBXxAwBUET8AQBXxAwBUET8AQBXxAwBUET8AQBXxAwBUET8AQBXxAwBUET8AQBXxAwBUET8AQBXxAwBUET8AQBXxAwBUET8AQBXxAwBUET8AQBXxAwBUET8AQBXxAwBUET8AQBXxAwBUET8AQBXxAwBUET8AQBXxAwBUET8AQBXxAwBUET8AQBXxAwBUET8AQBXxAwBUET8AQBXxAwBUET8AQBXxAwBUET8AQBXxAwBUET8AQBXxAwBUmd/2Aseb13Owefnd2PKetzf7aI6++otZr0/HnjKYxauP8rmLTS52j8eeMojPTXZzePw8k+3F2FMGszdZ5dH8cVZvz8aeMojl2VkOZu/Lq4vDsacM4qsm1zm/fpTV1dnYUwazPN3PR56/lPVXdXy+31se5+gDm1yc7Yw9ZVCbq82Nz0622+32poffPvl8nv7CPPPs3XjAi2Sdt/L29juTyXLsKYN5unsv//wbvz/Xs8XYUwYxzWnuzX40k8l67CmDOf9C8t/+1Cybq8nYUwbxgZ17+eTX/0D2ph339MX5eX7+Z/9LNpubv1G8aLZvzLP9tQ9mstcRP2fPd/Off/obstl0fID5Xdv1Nj/4Lft580O7X/HZW90Zx3lfTfgkySYnVeGTJOc7hzXhkySTXFSFT5IsH6cmfJLk/uKoJnySZLlcVoVPkuSVWU34JMnV1aIufJJkMptk5+hmf8DquTsAACJ+AIAy4gcAqCJ+AIAq4gcAqCJ+AIAq4gcAqCJ+AIAq4gcAqCJ+AIAq4gcAqCJ+AIAq4gcAqCJ+AIAq4gcAqCJ+AIAq4gcAqCJ+AIAq4gcAqCJ+AIAq4gcAqCJ+AIAq4gcAqCJ+AIAq4gcAqCJ+AIAq4gcAqCJ+AIAq4gcAqCJ+AIAq4gcAqCJ+AIAq4gcAqCJ+AIAq4gcAqCJ+AIAq4gcAqCJ+AIAq4gcAqCJ+AIAq4gcAqCJ+AIAq4gcAqCJ+AIAq4gcAqCJ+AIAq4gcAqCJ+AIAq4gcAqCJ+AIAq4gcAqCJ+AIAq4gcAqCJ+AIAq4gcAqCJ+AIAq4gcAqCJ+AIAq89teYPbl08xOL9+NLe9505xksX+c1d312FMGs9juZOfTZ3my3ht7yiCmmeY4d3Jnr+OeTpLN55IPvT5JjjdjTxnEq/urfOmjn0o2Hc/j58+e5QtPfjuXZ4uxpwxm+/5F1g/Ps7jcGXvKIE5P9nP+7DA7+8djTxnUdHqQ5OhGZ28VP7O3nuUDf/vfZrrqeBG5Ol7n7b/zdZnfOhlfHNOzvfzSz1xntbkYe8og7s6Sv/Tay1lMOkIgSaavrvNXfuZJJjuzsacMYjX5fP7Nnb+QbcljfHKZ/PAvz3K9HXvJcCYnycF3LTK5HnvJMHb3D/P1f/ovZ7nueA7/nnmuVj+Y5P4NTt7mn31+mZSET5Jc72+yLQqfJDm5WmS16XlCHUyXVeGTJNPjbSYdH5CTJOvJVU34JMn5KlXhkyST89SET5LMd3YznfW8Tv+e68yn5zc66Ts/AEAV8QMAVBE/AEAV8QMAVBE/AEAV8QMAVBE/AEAV8QMAVBE/AEAV8QMAVBE/AEAV8QMAVBE/AEAV8QMAVBE/AEAV8QMAVBE/AEAV8QMAVBE/AEAV8QMAVBE/AEAV8QMAVBE/AEAV8QMAVBE/AEAV8QMAVBE/AEAV8QMAVBE/AEAV8QMAVBE/AEAV8QMAVBE/AEAV8QMAVBE/AEAV8QMAVBE/AEAV8QMAVBE/AEAV8QMAVBE/AEAV8QMAVBE/AEAV8QMAVBE/AEAV8QMAVBE/AEAV8QMAVBE/AEAV8QMAVBE/AEAV8QMAVBE/AEAV8QMAVBE/AEAV8QMAVBE/AECV+W0vcLG7m+Vi593Y8p53uXed5TaZHKzHnjKY2dFBPvaRk0ynl2NPGcT9+fPMlkc5eXRv7CmDmb2zzubXtsmdJ2NPGcRyOslmvp87s9XYUwZx7zz58G8e5HJz65f7F8b0UTK/vJPlXsfPvL8+zvL61ayOdseeMqj5+vnNz97mH76YzPM//9g3ZjOd3eYyL4zt0Tqbr338LiTji+PVxWl+7rt/LLNpR/A9+uKr+fvf+uNZLTuCPkmuP/N23tr93qQkBu5ukk88u5NFtmNPGcTVxW6Of/Hj2W46XqeTZLm7m099/OPZpuNnvpwe5pe+9q9lMy96c0oy2azyfHJ9o7O3+p9aHu7XhE+SZG9TFT5Jsphf1IRPkjx7crcqfJJks3taEz5JcrC9rgmfJLleLarCJ0muF4tsi96blvt7deGTJNvpIpPdwxud9Z0fAKCK+AEAqogfAKCK+AEAqogfAKCK+AEAqogfAKCK+AEAqogfAKCK+AEAqogfAKCK+AEAqogfAKCK+AEAqogfAKCK+AEAqogfAKCK+AEAqogfAKCK+AEAqogfAKCK+AEAqogfAKCK+AEAqogfAKCK+AEAqogfAKCK+AEAqogfAKCK+AEAqogfAKCK+AEAqogfAKCK+AEAqogfAKCK+AEAqogfAKCK+AEAqogfAKCK+AEAqogfAKCK+AEAqogfAKCK+AEAqogfAKCK+AEAqogfAKCK+AEAqogfAKCK+AEAqogfAKCK+AEAqogfAKCK+AEAqogfAKCK+AEAqogfAKDK/LYXON1b52Jn+25sec/bHqyznF5lvns99pTBPNtb5gvXR7mXjsf4ej7P0Ud/PWfnq7GnDGZ69PmsD7e5nByNPWUQT9fT/Mbqa7K42h17yiAupwfJax/Kg/nV2FMGc3mwn+XkTq7ON2NPGcTx6V6efeZLWS7Ox54yqNM7x0lev9HZW8XPyWyVH/+mk6xnt7nKi2O7f5XVx34tmXWEQJLcXW/z/pOP5YtjDxnIs8PTbL/zT2Z3sx57ymCup8f55L2/m83k1p+FXgjb8zv5H//9u/MufPZ7Ibzxyiqf+rOfzX7R69bqdJvP/ug22/Vk7CmDWP/m0zz+lb+XpOeDeZIsF4vc/fiPJMcPvuKzt/qz19XRpCZ8kmS7c10VPklyuE0WY48Y0OXqIuui8EmSq8lBTfgkSZYHyabn572/WFeFT5KsL1ITPkmyyVnawidJdlarvHL+7EZnfecHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKvPbXuD4aDcHe4t3Y8t73vZolsX0lUw3V2NPGcz9SXJw9jS5OBt7yiCOny/z0eVruZ7e+qnxwnhpcz+/vTfNk6Oxlwxju7vO4YNHWZzujj1lEG/sXOWL82WWu5OxpwxmeXeb3/nIaXaXF2NPGcR13s7zxcuZX98de8qgZmenNz57q1f4g/1FvvfbviHzeccvkHbWV/m6J/uZzrZjTxnO85PMfupfZbJej71kGJP9fNPB92UymY29ZDCP7yWL79jmetHx5njw5Dp//rOfzPxBx+vWs/3kH/7RSdazjsc3SebXp/nQ1/xYpul43Zrlfh5M/0Umk52xpwxqslrmDx9ubnT2Vs/+u4d7NeGTJPPtdaYpCp8kOT/vCZ8kmexUhU+SnB2lJnySZPd8nvl1z+vW5U6qwidJZpuLmvBJkmnu1IVPkmwXO9m+/MqNzva8AgAARPwAAGXEDwBQRfwAAFXEDwBQRfwAAFXEDwBQRfwAAFXEDwBQRfwAAFXEDwBQRfwAAFXEDwBQRfwAAFXEDwBQRfwAAFXEDwBQRfwAAFXEDwBQRfwAAFXEDwBQRfwAAFXEDwBQRfwAAFXEDwBQRfwAAFXEDwBQRfwAAFXEDwBQRfwAAFXEDwBQRfwAAFXEDwBQRfwAAFXEDwBQRfwAAFXEDwBQRfwAAFXEDwBQRfwAAFXEDwBQRfwAAFXEDwBQRfwAAFXEDwBQRfwAAFXEDwBQRfwAAFXEDwBQRfwAAFXEDwBQRfwAAFXEDwBQRfwAAFXEDwBQRfwAAFXEDwBQZX7bC5w9neTqbPJubHnPO1/P86uf2+Z6dTb2lOE8usorl8lrd87HXjKM6TS7u29ld7Y39pLBLBbJwdn9PDm4N/aUQZzPd/PZ1y+zvX469pRBPDlKzp/Osni6M/aUwUxWz/P22UtZH+yOPWUQ87ycncPP5qX9Z2NPGdR2/UqSN2909lbx8+zJJP/ph46yue6In2fP1vmpn/zlrNfrsacM5rXj0/zEJ/5P3pp3/Mw78538kQ9PspkW/VJ0/9V8evb9WV13vDn+1r238hPf/h3ZZjX2lEEsfmeeN//qRzJd9tzTl/sH+V9/5nuymd368/0L4eWDt/JP/twfz2K2HHvKoLab3Wwmv57kq7/is7d6NqzOJzXhkySXl+dV4ZMk9w4uslMSPkmymM0zbQqfJKe7x1nNOsInSa63JzXhkyTzp7Oq8EmS1e5eTfgkydHuSV34JMlkepXZzuMbne16RgAA9cQPAFBF/AAAVcQPAFBF/AAAVcQPAFBF/AAAVcQPAFBF/AAAVcQPAFBF/AAAVcQPAFBF/AAAVcQPAFBF/AAAVcQPAFBF/AAAVcQPAFBF/AAAVcQPAFBF/AAAVcQPAFBF/AAAVcQPAFBF/AAAVcQPAFBF/AAAVcQPAFBF/AAAVcQPAFBF/AAAVcQPAFBF/AAAVcQPAFBF/AAAVcQPAFBF/AAAVcQPAFBF/AAAVcQPAFBF/AAAVcQPAFBF/AAAVcQPAFBF/AAAVcQPAFBF/AAAVcQPAFBF/AAAVcQPAFBF/AAAVcQPAFBF/AAAVcQPAFBF/AAAVcQPAFBF/AAAVcQPAFBF/AAAVW4VP0dHyXz+bk1579vbO8hsVvQDJ3l6vp/l9WzsGYNZra+z2WzGnjGo46vTLNbLsWcMZj65m0l2xp4xmOt762x2uu7pxdVlpuv12DMG8/zqblbrnnv6d00ne9mb37/R2cl2u93e5h9//Hib589vc4UXy9OnJ7m+vhh7xqDu7T3Ja/ea3hyT3d3DsWcM6su7L+XZ/itjzxjM2eatTHI69ozBbL+0zcHJ/tgzBvV8tpvc6XkeH+5+Ka8cPht7xqD25vdzuPPBG529dfwAALxIfOcHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKgifgCAKuIHAKjy/wAmD95oHcC2AAAAAABJRU5ErkJggg==\n"
          },
          "metadata": {}
        }
      ],
      "source": [
        "# 80 class\n",
        "# Generate a list of colors with a length equal to the number of labels\n",
        "colors = distinctipy.get_colors(len(class_names))\n",
        "\n",
        "# Make a copy of the color map in integer format\n",
        "int_colors = [tuple(int(c*255) for c in color) for color in colors]\n",
        "\n",
        "# Generate a color swatch to visualize the color map\n",
        "distinctipy.color_swatch(colors)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "axDrUtByegQh",
        "outputId": "dbcfd0ec-e571-43b0-ecf7-339ff050b8bf"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "The file ./KFOlCnqEu92Fr1MmEU9vAw.ttf already exists and overwrite is set to False.\n"
          ]
        }
      ],
      "source": [
        "# Set the name of the font file\n",
        "font_file = 'KFOlCnqEu92Fr1MmEU9vAw.ttf'\n",
        "\n",
        "# Download the font file\n",
        "download_file(f\"https://fonts.gstatic.com/s/roboto/v30/{font_file}\", \"./\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ocZ2F5fbeh9n"
      },
      "outputs": [],
      "source": [
        "draw_bboxes = partial(draw_bounding_boxes, fill=False, width=2, font=font_file, font_size=25)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o-FQXlsjIH3P"
      },
      "source": [
        "## Read video"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u-0T6ZD7IrAx"
      },
      "outputs": [],
      "source": [
        "# Open the video file\n",
        "cap = cv2.VideoCapture(video_path)\n",
        "\n",
        "# Get the video properties\n",
        "fps = int(cap.get(cv2.CAP_PROP_FPS))\n",
        "frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
        "frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
        "\n",
        "# Prepare a VideoWriter object to write the output video\n",
        "fourcc = cv2.VideoWriter_fourcc(*'mp4v')\n",
        "out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5uB4DEe2RxJT"
      },
      "outputs": [],
      "source": [
        "# from google.colab.patches import cv2_imshow\n",
        "\n",
        "# while cap.isOpened():\n",
        "#     ret, frame = cap.read()\n",
        "#     if not ret:\n",
        "#         break\n",
        "\n",
        "#     # Display the frame\n",
        "#     cv2_imshow(frame)\n",
        "\n",
        "#     # Wait for a key press\n",
        "#     if cv2.waitKey(1) & 0xFF == ord('q'):  # Press 'q' to quit\n",
        "#         break\n",
        "\n",
        "# # Release the video capture\n",
        "# cap.release()\n",
        "# cv2.destroyAllWindows()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y2rqMZ4A-cx3"
      },
      "source": [
        "## Predict"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Open the video file\n",
        "cap = cv2.VideoCapture(video_path)\n",
        "\n",
        "# Get the video properties\n",
        "fps = int(cap.get(cv2.CAP_PROP_FPS))\n",
        "frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
        "frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
        "\n",
        "# Prepare a VideoWriter object to write the output video\n",
        "fourcc = cv2.VideoWriter_fourcc(*'mp4v')\n",
        "out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))"
      ],
      "metadata": {
        "id": "ijLnZ40PiSXh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab.patches import cv2_imshow\n",
        "\n",
        "# Set the model to evaluation mode\n",
        "model.eval()\n",
        "\n",
        "# Ensure the model and input data are on the same device\n",
        "model.to(device)\n",
        "\n",
        "while cap.isOpened():\n",
        "    ret, frame = cap.read()\n",
        "    if not ret:\n",
        "        break\n",
        "\n",
        "    # Convert the frame to RGB color space\n",
        "    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)\n",
        "\n",
        "    # Convert the frame to a tensor and move it to the correct device\n",
        "    input_tensor = transforms.ToTensor()(frame_rgb).unsqueeze(0).to(device)\n",
        "\n",
        "    # Make a prediction with the model\n",
        "    with torch.no_grad():\n",
        "        model_output = model(input_tensor)\n",
        "\n",
        "    # Set the confidence threshold\n",
        "    threshold = 0.5\n",
        "\n",
        "    # Filter the output based on the confidence threshold\n",
        "    scores_mask = model_output[0]['scores'] > threshold\n",
        "\n",
        "    # Scale the predicted bounding boxes\n",
        "    pred_bboxes = model_output[0]['boxes'][scores_mask]\n",
        "\n",
        "    # Get the class names for the predicted label indices\n",
        "    pred_labels = [class_names[int(label)] for label in model_output[0]['labels'][scores_mask]]\n",
        "\n",
        "    # Extract the confidence scores\n",
        "    pred_scores = model_output[0]['scores'][scores_mask]\n",
        "\n",
        "    # Resize and stack the predicted segmentation masks\n",
        "    pred_masks = model_output[0]['masks'][scores_mask].unsqueeze(0)\n",
        "    pred_masks = torch.stack([torch.where(mask >= threshold, 1, 0) for mask in pred_masks])\n",
        "\n",
        "    # Annotate the frame with the predicted segmentation masks\n",
        "    annotated_frame = frame_rgb.copy()\n",
        "    for mask, label, bbox, score in zip(pred_masks, pred_labels, pred_bboxes, pred_scores):\n",
        "        annotated_frame = cv2.putText(annotated_frame, f\"{label}: {score*100:.2f}%\", (int(bbox[0]), int(bbox[1])), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)\n",
        "        annotated_frame = cv2.rectangle(annotated_frame, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), (0, 255, 0), 2)\n",
        "\n",
        "    # Display the annotated frame\n",
        "    # cv2_imshow(annotated_frame)\n",
        "\n",
        "    # Write the annotated frame to the video\n",
        "    out.write(annotated_frame)\n",
        "\n",
        "# Release the video capture\n",
        "out.release()\n",
        "cap.release()\n",
        "cv2.destroyAllWindows()\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7QCExcgRUrFy",
        "outputId": "43d33ab0-68f2-4e82-ab79-79379e398856"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/torchvision/transforms/v2/_deprecated.py:42: UserWarning: The transform `ToTensor()` is deprecated and will be removed in a future release. Instead, please use `v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True)])`.Output is equivalent up to float precision.\n",
            "  warnings.warn(\n"
          ]
        }
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}