{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "976ee38f",
   "metadata": {},
   "source": [
    "## Question 7 - YOLOv5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9b05cf93",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "\n",
    "from pathlib import Path\n",
    "import xml.etree.ElementTree as ET\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import Dataset\n",
    "from torchvision import transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cb4aa690",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Windows Fix\n",
    "# import pathlib\n",
    "# pathlib.PosixPath = pathlib.WindowsPath"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1a045ed5",
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_NAME = 'data/yolov5.pt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "190ad178",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define transform\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize((640, 640)),\n",
    "    transforms.ToTensor(),\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0011f88c",
   "metadata": {},
   "outputs": [],
   "source": [
    "device='cuda' if torch.cuda.is_available() else 'cpu'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3b7d1c5",
   "metadata": {},
   "source": [
    "## Convert data format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "76320268",
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_xml_to_yolo(root_dir):\n",
    "    # Directories for the YOLO formatted annotations\n",
    "    os.makedirs(os.path.join(root_dir, 'labels'), exist_ok=True)\n",
    "    images_dir = os.path.join(root_dir, 'images')\n",
    "    os.makedirs(images_dir, exist_ok=True)\n",
    "\n",
    "    # Class mapping to index\n",
    "    class_mapping = {'with_mask': 0, 'mask_weared_incorrect': 1, 'without_mask': 2}\n",
    "\n",
    "    # Process each XML file\n",
    "    for filename in os.listdir(root_dir):\n",
    "        if not filename.endswith('.xml'):\n",
    "            continue\n",
    "\n",
    "        # Parse the XML File\n",
    "        tree = ET.parse(os.path.join(root_dir, filename))\n",
    "        root = tree.getroot()\n",
    "        image_name = filename.replace('.xml', '')\n",
    "        image_path = None\n",
    "\n",
    "        # Find corresponding image file and convert to RGB\n",
    "        for ext in ['.jpg', '.png']:\n",
    "            if os.path.exists(os.path.join(root_dir, image_name + ext)):\n",
    "                image_path = os.path.join(root_dir, image_name + ext)\n",
    "                break\n",
    "\n",
    "        if image_path:\n",
    "            # Move image to 'images' folder and convert to RGB\n",
    "            image = Image.open(image_path).convert(\"RGB\")\n",
    "            rgb_image_path = os.path.join(images_dir, image_name + '.png')  # Saving as PNG\n",
    "            image.save(rgb_image_path)\n",
    "            w, h = image.size\n",
    "\n",
    "            # Create a corresponding YOLO format file\n",
    "            with open(os.path.join(root_dir, 'labels', filename.replace('.xml', '.txt')), 'w') as f:\n",
    "                for obj in root.findall('object'):\n",
    "                    class_name = obj.find('name').text\n",
    "                    if class_name not in class_mapping:\n",
    "                        continue  # Skip unknown classes\n",
    "\n",
    "                    # Convert XML box to YOLO format\n",
    "                    bbox = obj.find('bndbox')\n",
    "                    xmin = int(bbox.find('xmin').text)\n",
    "                    ymin = int(bbox.find('ymin').text)\n",
    "                    xmax = int(bbox.find('xmax').text)\n",
    "                    ymax = int(bbox.find('ymax').text)\n",
    "                    x_center = ((xmin + xmax) / 2) / w\n",
    "                    y_center = ((ymin + ymax) / 2) / h\n",
    "                    width = (xmax - xmin) / w\n",
    "                    height = (ymax - ymin) / h\n",
    "\n",
    "                    # Write to file\n",
    "                    f.write(f\"{class_mapping[class_name]} {x_center} {y_center} {width} {height}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "fd415c2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert both train and validation datasets\n",
    "# convert_xml_to_yolo('data/MaskedFace/train')\n",
    "# convert_xml_to_yolo('data/MaskedFace/val')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53e95723",
   "metadata": {},
   "source": [
    "## Get the ground truth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ee0db89f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_true_labels_counts(label_files_dir, num_classes):\n",
    "    # List all label files\n",
    "    label_files = list(Path(label_files_dir).rglob('*.txt'))\n",
    "    label_files.sort()\n",
    "    counts = np.zeros((len(label_files), num_classes), dtype=int)\n",
    "    filenames = []\n",
    "    bounding_boxes = []\n",
    "\n",
    "    for i, file_path in enumerate(label_files):\n",
    "        with open(file_path) as f:\n",
    "            box = []\n",
    "            for line in f:\n",
    "                class_id = int(line.split()[0])\n",
    "                x_min = float(line.split()[1])\n",
    "                y_min = float(line.split()[2])\n",
    "                x_max = float(line.split()[3])\n",
    "                y_max = float(line.split()[4])\n",
    "                counts[i, class_id] += 1\n",
    "                box.append([x_min, y_min, x_max, y_max])\n",
    "        filenames.append(file_path.name)\n",
    "        bounding_boxes.append(box)\n",
    "\n",
    "    return bounding_boxes, counts, filenames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "742f36b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example usage\n",
    "label_files_dir = 'data/MaskedFace/val/labels'\n",
    "num_classes = 3\n",
    "\n",
    "gt_boxes, gt_counts, gt_filenames = extract_true_labels_counts(label_files_dir, num_classes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "8537e68c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ground Truth counts shape: (85, 3)\n"
     ]
    }
   ],
   "source": [
    "print(\"Ground Truth counts shape:\", gt_counts.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7e238b6",
   "metadata": {},
   "source": [
    "## Load the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "9db12843",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append(\"yolov5\")\n",
    "from models.experimental import attempt_load\n",
    "from utils.general import non_max_suppression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "41e74112",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MaskedFaceDataset(Dataset):\n",
    "    def __init__(self, root, transform=None):\n",
    "        super().__init__()\n",
    "        self.img_files = list(Path(root).rglob('*.png'))\n",
    "        self.img_files.sort()\n",
    "        self.transform = transform\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        img_path = str(self.img_files[idx])\n",
    "        img = Image.open(img_path).convert('RGB')\n",
    "        if self.transform:\n",
    "            img = self.transform(img)\n",
    "            filename=os.path.basename(img_path)\n",
    "        return img, filename  # Return both the image and its path for visualization\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.img_files)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "86d7aa7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "val_dataset = MaskedFaceDataset(root='data/MaskedFace/val/images', transform=transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "59de8e25",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Fusing layers... \n",
      "Model summary: 157 layers, 7018216 parameters, 0 gradients, 15.8 GFLOPs\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "DetectionModel(\n",
       "  (model): Sequential(\n",
       "    (0): Conv(\n",
       "      (conv): Conv2d(3, 32, kernel_size=(6, 6), stride=(2, 2), padding=(2, 2))\n",
       "      (act): SiLU(inplace=True)\n",
       "    )\n",
       "    (1): Conv(\n",
       "      (conv): Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
       "      (act): SiLU(inplace=True)\n",
       "    )\n",
       "    (2): C3(\n",
       "      (cv1): Conv(\n",
       "        (conv): Conv2d(64, 32, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (act): SiLU(inplace=True)\n",
       "      )\n",
       "      (cv2): Conv(\n",
       "        (conv): Conv2d(64, 32, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (act): SiLU(inplace=True)\n",
       "      )\n",
       "      (cv3): Conv(\n",
       "        (conv): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (act): SiLU(inplace=True)\n",
       "      )\n",
       "      (m): Sequential(\n",
       "        (0): Bottleneck(\n",
       "          (cv1): Conv(\n",
       "            (conv): Conv2d(32, 32, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (act): SiLU(inplace=True)\n",
       "          )\n",
       "          (cv2): Conv(\n",
       "            (conv): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            (act): SiLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (3): Conv(\n",
       "      (conv): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
       "      (act): SiLU(inplace=True)\n",
       "    )\n",
       "    (4): C3(\n",
       "      (cv1): Conv(\n",
       "        (conv): Conv2d(128, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (act): SiLU(inplace=True)\n",
       "      )\n",
       "      (cv2): Conv(\n",
       "        (conv): Conv2d(128, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (act): SiLU(inplace=True)\n",
       "      )\n",
       "      (cv3): Conv(\n",
       "        (conv): Conv2d(128, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (act): SiLU(inplace=True)\n",
       "      )\n",
       "      (m): Sequential(\n",
       "        (0): Bottleneck(\n",
       "          (cv1): Conv(\n",
       "            (conv): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (act): SiLU(inplace=True)\n",
       "          )\n",
       "          (cv2): Conv(\n",
       "            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            (act): SiLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "        (1): Bottleneck(\n",
       "          (cv1): Conv(\n",
       "            (conv): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (act): SiLU(inplace=True)\n",
       "          )\n",
       "          (cv2): Conv(\n",
       "            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            (act): SiLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (5): Conv(\n",
       "      (conv): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
       "      (act): SiLU(inplace=True)\n",
       "    )\n",
       "    (6): C3(\n",
       "      (cv1): Conv(\n",
       "        (conv): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (act): SiLU(inplace=True)\n",
       "      )\n",
       "      (cv2): Conv(\n",
       "        (conv): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (act): SiLU(inplace=True)\n",
       "      )\n",
       "      (cv3): Conv(\n",
       "        (conv): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (act): SiLU(inplace=True)\n",
       "      )\n",
       "      (m): Sequential(\n",
       "        (0): Bottleneck(\n",
       "          (cv1): Conv(\n",
       "            (conv): Conv2d(128, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (act): SiLU(inplace=True)\n",
       "          )\n",
       "          (cv2): Conv(\n",
       "            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            (act): SiLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "        (1): Bottleneck(\n",
       "          (cv1): Conv(\n",
       "            (conv): Conv2d(128, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (act): SiLU(inplace=True)\n",
       "          )\n",
       "          (cv2): Conv(\n",
       "            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            (act): SiLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "        (2): Bottleneck(\n",
       "          (cv1): Conv(\n",
       "            (conv): Conv2d(128, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (act): SiLU(inplace=True)\n",
       "          )\n",
       "          (cv2): Conv(\n",
       "            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            (act): SiLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (7): Conv(\n",
       "      (conv): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
       "      (act): SiLU(inplace=True)\n",
       "    )\n",
       "    (8): C3(\n",
       "      (cv1): Conv(\n",
       "        (conv): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (act): SiLU(inplace=True)\n",
       "      )\n",
       "      (cv2): Conv(\n",
       "        (conv): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (act): SiLU(inplace=True)\n",
       "      )\n",
       "      (cv3): Conv(\n",
       "        (conv): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (act): SiLU(inplace=True)\n",
       "      )\n",
       "      (m): Sequential(\n",
       "        (0): Bottleneck(\n",
       "          (cv1): Conv(\n",
       "            (conv): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (act): SiLU(inplace=True)\n",
       "          )\n",
       "          (cv2): Conv(\n",
       "            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            (act): SiLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (9): SPPF(\n",
       "      (cv1): Conv(\n",
       "        (conv): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (act): SiLU(inplace=True)\n",
       "      )\n",
       "      (cv2): Conv(\n",
       "        (conv): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (act): SiLU(inplace=True)\n",
       "      )\n",
       "      (m): MaxPool2d(kernel_size=5, stride=1, padding=2, dilation=1, ceil_mode=False)\n",
       "    )\n",
       "    (10): Conv(\n",
       "      (conv): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "      (act): SiLU(inplace=True)\n",
       "    )\n",
       "    (11): Upsample(scale_factor=2.0, mode=nearest)\n",
       "    (12): Concat()\n",
       "    (13): C3(\n",
       "      (cv1): Conv(\n",
       "        (conv): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (act): SiLU(inplace=True)\n",
       "      )\n",
       "      (cv2): Conv(\n",
       "        (conv): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (act): SiLU(inplace=True)\n",
       "      )\n",
       "      (cv3): Conv(\n",
       "        (conv): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (act): SiLU(inplace=True)\n",
       "      )\n",
       "      (m): Sequential(\n",
       "        (0): Bottleneck(\n",
       "          (cv1): Conv(\n",
       "            (conv): Conv2d(128, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (act): SiLU(inplace=True)\n",
       "          )\n",
       "          (cv2): Conv(\n",
       "            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            (act): SiLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (14): Conv(\n",
       "      (conv): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "      (act): SiLU(inplace=True)\n",
       "    )\n",
       "    (15): Upsample(scale_factor=2.0, mode=nearest)\n",
       "    (16): Concat()\n",
       "    (17): C3(\n",
       "      (cv1): Conv(\n",
       "        (conv): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (act): SiLU(inplace=True)\n",
       "      )\n",
       "      (cv2): Conv(\n",
       "        (conv): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (act): SiLU(inplace=True)\n",
       "      )\n",
       "      (cv3): Conv(\n",
       "        (conv): Conv2d(128, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (act): SiLU(inplace=True)\n",
       "      )\n",
       "      (m): Sequential(\n",
       "        (0): Bottleneck(\n",
       "          (cv1): Conv(\n",
       "            (conv): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (act): SiLU(inplace=True)\n",
       "          )\n",
       "          (cv2): Conv(\n",
       "            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            (act): SiLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (18): Conv(\n",
       "      (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
       "      (act): SiLU(inplace=True)\n",
       "    )\n",
       "    (19): Concat()\n",
       "    (20): C3(\n",
       "      (cv1): Conv(\n",
       "        (conv): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (act): SiLU(inplace=True)\n",
       "      )\n",
       "      (cv2): Conv(\n",
       "        (conv): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (act): SiLU(inplace=True)\n",
       "      )\n",
       "      (cv3): Conv(\n",
       "        (conv): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (act): SiLU(inplace=True)\n",
       "      )\n",
       "      (m): Sequential(\n",
       "        (0): Bottleneck(\n",
       "          (cv1): Conv(\n",
       "            (conv): Conv2d(128, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (act): SiLU(inplace=True)\n",
       "          )\n",
       "          (cv2): Conv(\n",
       "            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            (act): SiLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (21): Conv(\n",
       "      (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
       "      (act): SiLU(inplace=True)\n",
       "    )\n",
       "    (22): Concat()\n",
       "    (23): C3(\n",
       "      (cv1): Conv(\n",
       "        (conv): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (act): SiLU(inplace=True)\n",
       "      )\n",
       "      (cv2): Conv(\n",
       "        (conv): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (act): SiLU(inplace=True)\n",
       "      )\n",
       "      (cv3): Conv(\n",
       "        (conv): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (act): SiLU(inplace=True)\n",
       "      )\n",
       "      (m): Sequential(\n",
       "        (0): Bottleneck(\n",
       "          (cv1): Conv(\n",
       "            (conv): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "            (act): SiLU(inplace=True)\n",
       "          )\n",
       "          (cv2): Conv(\n",
       "            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "            (act): SiLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (24): Detect(\n",
       "      (m): ModuleList(\n",
       "        (0): Conv2d(128, 24, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (1): Conv2d(256, 24, kernel_size=(1, 1), stride=(1, 1))\n",
       "        (2): Conv2d(512, 24, kernel_size=(1, 1), stride=(1, 1))\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = attempt_load(MODEL_NAME, device=device)\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3113518b",
   "metadata": {},
   "source": [
    "## Test an image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "601f0cbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "c69a8766",
   "metadata": {},
   "outputs": [],
   "source": [
    "index = 20\n",
    "img, img_file = val_dataset[index]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d191f4c7",
   "metadata": {},
   "source": [
    "Ground truth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "69a98357",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 288888 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "cv_img = (img.permute(1, 2, 0).numpy() * 255.0).astype(np.uint8).copy()\n",
    "height, width, _ = cv_img.shape\n",
    "\n",
    "for box in gt_boxes[index]:\n",
    "    xc, yc, w, h = box\n",
    "    x_min = int((xc - w / 2) * width)\n",
    "    x_max = int((xc + w / 2) * width)\n",
    "    y_min = int((yc - h / 2) * height)\n",
    "    y_max = int((yc + h / 2) * height)\n",
    "    cv2.rectangle( cv_img, \n",
    "                  (x_min, y_min), \n",
    "                  (x_max, y_max), (255, 255, 255), 2)\n",
    "\n",
    "plt.imshow(cv_img)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "6587fb50",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 289272 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    pred = model(img.unsqueeze(0).to(device))\n",
    "    pred = non_max_suppression(pred, 0.25, 0.45, classes=None, agnostic=False)\n",
    "\n",
    "cv_img = (img.permute(1, 2, 0).numpy() * 255.0).astype(np.uint8).copy()\n",
    "\n",
    "for det in pred[0]:\n",
    "#     print(det[0:2], det[2:4], det[4], det[5])\n",
    "    cv2.rectangle( cv_img, (int(det[0]), int(det[1])), (int(det[2]), int(det[3])), (255, 255, 255), 2)\n",
    "\n",
    "plt.imshow(cv_img)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7c43d8b",
   "metadata": {},
   "source": [
    "## Test on the validation set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "7604702c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "def count_masks(model, dataset):\n",
    "\n",
    "    predicted_counts = np.zeros((len(dataset), 3), dtype=int)\n",
    "    filenames = []\n",
    "\n",
    "    for idx, (img, filename) in enumerate(tqdm(dataset)):\n",
    "        img = img.unsqueeze(0).to(device)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            pred = model(img)\n",
    "            pred = non_max_suppression(pred, 0.25, 0.45, classes=None, agnostic=False)\n",
    "\n",
    "        if len(pred) > 0 and pred[0] is not None:\n",
    "            for det in pred[0]:\n",
    "                if det is None or len(det) == 0:\n",
    "                    continue\n",
    "                if det.dim() == 1:\n",
    "                    det = det.unsqueeze(0)\n",
    "                cls_ids = det[:, -1].cpu().numpy()\n",
    "                for cls_id in cls_ids:\n",
    "                    predicted_counts[idx, int(cls_id)] += 1\n",
    "\n",
    "    return predicted_counts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "ade29625",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 85/85 [00:01<00:00, 75.81it/s]\n"
     ]
    }
   ],
   "source": [
    "predicted_counts = count_masks(model, val_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d92cdd7",
   "metadata": {},
   "source": [
    "## MAPE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "9ee27bdc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_mape(prediction, truth):\n",
    "    mape = np.mean( np.abs(truth - prediction) / np.maximum(truth, np.ones_like(truth)) ) * 100\n",
    "    return mape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "f6dfb0d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "MAPE = compute_mape(predicted_counts, gt_counts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "b5954e63",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "28.946963062185922\n"
     ]
    }
   ],
   "source": [
    "print(MAPE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae09c24e",
   "metadata": {},
   "source": [
    "## Final Score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "bdb0b01a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Score:  12.5\n"
     ]
    }
   ],
   "source": [
    "if MAPE <= 10:\n",
    "    print(\"Score: \", 25*1.0)\n",
    "elif MAPE <= 15:\n",
    "    print(\"Score: \", 25*0.875)\n",
    "elif MAPE <= 20:\n",
    "    print(\"Score: \", 25*0.75)\n",
    "elif MAPE <= 25:\n",
    "    print(\"Score: \", 25*0.625)\n",
    "elif MAPE <= 30:\n",
    "    print(\"Score: \", 25*0.5)\n",
    "else:\n",
    "    print(\"Score: \", 0)  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb2a99e2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "what",
   "language": "python",
   "name": "what"
  },
  "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
