{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6b31398c",
   "metadata": {},
   "source": [
    "## Question 7 - FasterRCNN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0f36ed82",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "\n",
    "import cv2\n",
    "import torch\n",
    "import torchvision\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "\n",
    "import pandas as pd\n",
    "from pathlib import Path\n",
    "\n",
    "from tqdm import tqdm\n",
    "from collections import Counter\n",
    "from xml.etree import ElementTree as ET\n",
    "\n",
    "from torchvision import transforms, models\n",
    "from torch.utils.data import Dataset, DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4422b46e",
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_NAME = 'data/faster_rcnn_fold_5_MSKmodel.pth'\n",
    "\n",
    "class_dict = {\"with_mask\": 2, \n",
    "              \"without_mask\": 1, \n",
    "              \"mask_weared_incorrect\": 3}\n",
    "\n",
    "class_dct_inv  = {2: 'with_mask', \n",
    "                  1: \"without_mask\", \n",
    "                  3: \"mask_weared_incorrect\"}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "740b253c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import albumentations as A #note installing this changes your version of OpenCV and you loose sift\n",
    "from albumentations.pytorch import ToTensorV2\n",
    "\n",
    "bbox_params = A.BboxParams(\n",
    "        format = 'pascal_voc',\n",
    "        label_fields = ['labels'],\n",
    "        min_visibility = 0.3\n",
    ")\n",
    "\n",
    "test_transform = A.Compose([\n",
    "        A.Resize(height=480, width=480),\n",
    "#         A.Normalize(\n",
    "#             mean=[0.485, 0.456, 0.406],\n",
    "#             std=[0.229, 0.224, 0.225]),\n",
    "        ToTensorV2(p=1.0)],\n",
    "        bbox_params = bbox_params\n",
    ")\n",
    "\n",
    "# bbox_params = A.BboxParams(\n",
    "#     format='pascal_voc',\n",
    "#     min_area=5,  # 5\n",
    "#     min_visibility=0.9,  # 0.9\n",
    "#     #     check_each_transform =True,\n",
    "#     label_fields=['labels']\n",
    "# )\n",
    "\n",
    "# test_transform = A.Compose([\n",
    "#     A.Resize(height=446, width=446),\n",
    "#     ToTensorV2()],\n",
    "#     bbox_params=bbox_params,\n",
    "#     is_check_shapes=False\n",
    "# )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "871b113a",
   "metadata": {},
   "source": [
    "## Load the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6b80c289",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.models.detection.faster_rcnn import FastRCNNPredictor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7d9a0849",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_model_detection(num_classes):\n",
    "    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)\n",
    "    in_features = model.roi_heads.box_predictor.cls_score.in_features\n",
    "    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a71caa3f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "D:\\Anaconda3\\envs\\what\\lib\\site-packages\\torchvision\\models\\_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.\n",
      "  warnings.warn(\n",
      "D:\\Anaconda3\\envs\\what\\lib\\site-packages\\torchvision\\models\\_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=FasterRCNN_ResNet50_FPN_Weights.COCO_V1`. You can also use `weights=FasterRCNN_ResNet50_FPN_Weights.DEFAULT` to get the most up-to-date weights.\n",
      "  warnings.warn(msg)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "FasterRCNN(\n",
       "  (transform): GeneralizedRCNNTransform(\n",
       "      Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
       "      Resize(min_size=(800,), max_size=1333, mode='bilinear')\n",
       "  )\n",
       "  (backbone): BackboneWithFPN(\n",
       "    (body): IntermediateLayerGetter(\n",
       "      (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
       "      (bn1): FrozenBatchNorm2d(64, eps=0.0)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
       "      (layer1): Sequential(\n",
       "        (0): Bottleneck(\n",
       "          (conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(64, eps=0.0)\n",
       "          (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(64, eps=0.0)\n",
       "          (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (downsample): Sequential(\n",
       "            (0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          )\n",
       "        )\n",
       "        (1): Bottleneck(\n",
       "          (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(64, eps=0.0)\n",
       "          (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(64, eps=0.0)\n",
       "          (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "        )\n",
       "        (2): Bottleneck(\n",
       "          (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(64, eps=0.0)\n",
       "          (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(64, eps=0.0)\n",
       "          (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "        )\n",
       "      )\n",
       "      (layer2): Sequential(\n",
       "        (0): Bottleneck(\n",
       "          (conv1): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(128, eps=0.0)\n",
       "          (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(128, eps=0.0)\n",
       "          (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(512, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (downsample): Sequential(\n",
       "            (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "            (1): FrozenBatchNorm2d(512, eps=0.0)\n",
       "          )\n",
       "        )\n",
       "        (1): Bottleneck(\n",
       "          (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(128, eps=0.0)\n",
       "          (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(128, eps=0.0)\n",
       "          (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(512, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "        )\n",
       "        (2): Bottleneck(\n",
       "          (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(128, eps=0.0)\n",
       "          (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(128, eps=0.0)\n",
       "          (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(512, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "        )\n",
       "        (3): Bottleneck(\n",
       "          (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(128, eps=0.0)\n",
       "          (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(128, eps=0.0)\n",
       "          (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(512, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "        )\n",
       "      )\n",
       "      (layer3): Sequential(\n",
       "        (0): Bottleneck(\n",
       "          (conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(1024, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (downsample): Sequential(\n",
       "            (0): Conv2d(512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "            (1): FrozenBatchNorm2d(1024, eps=0.0)\n",
       "          )\n",
       "        )\n",
       "        (1): Bottleneck(\n",
       "          (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(1024, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "        )\n",
       "        (2): Bottleneck(\n",
       "          (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(1024, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "        )\n",
       "        (3): Bottleneck(\n",
       "          (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(1024, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "        )\n",
       "        (4): Bottleneck(\n",
       "          (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(1024, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "        )\n",
       "        (5): Bottleneck(\n",
       "          (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(256, eps=0.0)\n",
       "          (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(1024, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "        )\n",
       "      )\n",
       "      (layer4): Sequential(\n",
       "        (0): Bottleneck(\n",
       "          (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(512, eps=0.0)\n",
       "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(512, eps=0.0)\n",
       "          (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(2048, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (downsample): Sequential(\n",
       "            (0): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "            (1): FrozenBatchNorm2d(2048, eps=0.0)\n",
       "          )\n",
       "        )\n",
       "        (1): Bottleneck(\n",
       "          (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(512, eps=0.0)\n",
       "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(512, eps=0.0)\n",
       "          (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(2048, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "        )\n",
       "        (2): Bottleneck(\n",
       "          (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn1): FrozenBatchNorm2d(512, eps=0.0)\n",
       "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): FrozenBatchNorm2d(512, eps=0.0)\n",
       "          (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn3): FrozenBatchNorm2d(2048, eps=0.0)\n",
       "          (relu): ReLU(inplace=True)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (fpn): FeaturePyramidNetwork(\n",
       "      (inner_blocks): ModuleList(\n",
       "        (0): Conv2dNormActivation(\n",
       "          (0): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "        )\n",
       "        (1): Conv2dNormActivation(\n",
       "          (0): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "        )\n",
       "        (2): Conv2dNormActivation(\n",
       "          (0): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "        )\n",
       "        (3): Conv2dNormActivation(\n",
       "          (0): Conv2d(2048, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "        )\n",
       "      )\n",
       "      (layer_blocks): ModuleList(\n",
       "        (0-3): 4 x Conv2dNormActivation(\n",
       "          (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "        )\n",
       "      )\n",
       "      (extra_blocks): LastLevelMaxPool()\n",
       "    )\n",
       "  )\n",
       "  (rpn): RegionProposalNetwork(\n",
       "    (anchor_generator): AnchorGenerator()\n",
       "    (head): RPNHead(\n",
       "      (conv): Sequential(\n",
       "        (0): Conv2dNormActivation(\n",
       "          (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "          (1): ReLU(inplace=True)\n",
       "        )\n",
       "      )\n",
       "      (cls_logits): Conv2d(256, 3, kernel_size=(1, 1), stride=(1, 1))\n",
       "      (bbox_pred): Conv2d(256, 12, kernel_size=(1, 1), stride=(1, 1))\n",
       "    )\n",
       "  )\n",
       "  (roi_heads): RoIHeads(\n",
       "    (box_roi_pool): MultiScaleRoIAlign(featmap_names=['0', '1', '2', '3'], output_size=(7, 7), sampling_ratio=2)\n",
       "    (box_head): TwoMLPHead(\n",
       "      (fc6): Linear(in_features=12544, out_features=1024, bias=True)\n",
       "      (fc7): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "    )\n",
       "    (box_predictor): FastRCNNPredictor(\n",
       "      (cls_score): Linear(in_features=1024, out_features=4, bias=True)\n",
       "      (bbox_pred): Linear(in_features=1024, out_features=16, bias=True)\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device='cuda' if torch.cuda.is_available() else 'cpu'\n",
    "model = get_model_detection(4)\n",
    "\n",
    "model.load_state_dict(torch.load(MODEL_NAME, map_location=torch.device('cpu')), False)\n",
    "model.eval()\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82856a39",
   "metadata": {},
   "source": [
    "## Load the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "36c950bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from bs4 import BeautifulSoup\n",
    "\n",
    "class MaskedFaceTestDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, path_to_imgs, path_to_xml_files, transform=None):\n",
    "        self.path_to_imgs = path_to_imgs\n",
    "        self.path_to_xml_files = path_to_xml_files\n",
    "        self.transform = transform\n",
    "\n",
    "        self.file_names = [file for file in os.listdir(self.path_to_imgs) if file.endswith(\".png\")]\n",
    "\n",
    "        self.file_paths = [os.path.join(self.path_to_imgs, file_name) for file_name in self.file_names]\n",
    "\n",
    "    def __getitem__(self, index: int):\n",
    "        path_to_img = self.file_paths[index]\n",
    "        path = Path(path_to_img)\n",
    "        fname = path.stem\n",
    "\n",
    "        self.name = self.file_names[index]\n",
    "\n",
    "        image = Image.open(path_to_img).convert('RGB')\n",
    "        image = np.array(image) / 255.0\n",
    "\n",
    "        target = self.__generate_target(self.path_to_xml_files, self.name)\n",
    "\n",
    "        if self.transform:\n",
    "            augmentations = self.transform(image=image, bboxes=target['boxes'], labels=target['labels'])\n",
    "\n",
    "            image = augmentations['image']\n",
    "\n",
    "            target['boxes'] = torch.tensor(augmentations['bboxes'])\n",
    "\n",
    "            target['labels'] = torch.tensor(augmentations['labels'])\n",
    "\n",
    "        return image, target, fname\n",
    "\n",
    "    def __generate_target(self, path_to_xml, name_of_img):\n",
    "        path_to_annot = os.path.join(path_to_xml, name_of_img.replace('png', 'xml'))\n",
    "\n",
    "        boxes_lst = []\n",
    "        labels_lst = []\n",
    "        clases_lst = []\n",
    "        target = {}\n",
    "\n",
    "        with open(path_to_annot, 'r') as file:\n",
    "            xml_content = file.read()\n",
    "\n",
    "        soup = BeautifulSoup(xml_content, 'xml')\n",
    "\n",
    "        objects = soup.find_all('object')\n",
    "\n",
    "        for obj in objects:\n",
    "            name = obj.find('name').text\n",
    "            xmin = int(float(obj.find('xmin').text)) - 1\n",
    "            ymin = int(float(obj.find('ymin').text)) - 1\n",
    "            xmax = int(float(obj.find('xmax').text)) - 1\n",
    "            ymax = int(float(obj.find('ymax').text)) - 1\n",
    "\n",
    "            labels_lst.append(name)\n",
    "            boxes_lst.append([xmin, ymin, xmax, ymax])\n",
    "            clases_lst.append(class_dict[name])\n",
    "\n",
    "        target[\"boxes\"] = torch.tensor(boxes_lst, dtype=torch.float32)\n",
    "        target[\"labels\"] = torch.tensor(clases_lst, dtype=torch.int64)\n",
    "\n",
    "        return target\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.file_paths)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5d486c56",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import xmltodict\n",
    "from torch.utils.data.sampler import SubsetRandomSampler\n",
    "\n",
    "def collate_fn(batch):\n",
    "    return tuple(zip(*batch))\n",
    "\n",
    "test_ds = MaskedFaceTestDataset(path_to_imgs = './data/MaskedFace/val/',\n",
    "                               path_to_xml_files = './data/MaskedFace/val/',\n",
    "                               transform = test_transform)\n",
    "\n",
    "test_sampler = SubsetRandomSampler( list(range(0, len(test_ds))) )\n",
    "\n",
    "test_loader = DataLoader(test_ds, batch_size=16, num_workers=0, pin_memory=True,\n",
    "                             sampler=test_sampler, collate_fn=collate_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5060b3a",
   "metadata": {},
   "source": [
    "## Test an image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "837351e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "index = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f0899a08",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 335516 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "image, target, fname = test_ds[0]\n",
    "\n",
    "cv_image = np.transpose(image, (1, 2, 0)).numpy()\n",
    "\n",
    "for box, label in zip(target['boxes'], target['labels']):\n",
    "    cv2.rectangle(cv_image, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (255, 255, 0))\n",
    "    cv2.putText(cv_image, \n",
    "                class_dct_inv[label.item()], \n",
    "                (int(box[0])-10, int(box[1])-10),\n",
    "                cv2.FONT_HERSHEY_DUPLEX,\n",
    "                0.5,\n",
    "                (255, 255, 0), \n",
    "               )\n",
    "\n",
    "plt.imshow(cv_image)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "1f31baf5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 335516 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    prediction = model(image.unsqueeze(0).to(device, dtype=torch.float))[0]\n",
    "\n",
    "cv_image = np.transpose(image, (1, 2, 0)).numpy()\n",
    "\n",
    "for box, label, score in zip(prediction['boxes'], prediction['labels'], prediction['scores']):\n",
    "    if score >= 0.5:\n",
    "        cv2.rectangle(cv_image, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (255, 255, 0))\n",
    "        cv2.putText(cv_image, \n",
    "                    class_dct_inv[label.item()], \n",
    "                    (int(box[0])-10, int(box[1])-10),\n",
    "                    cv2.FONT_HERSHEY_DUPLEX,\n",
    "                    0.5,\n",
    "                    (255, 255, 0), \n",
    "                   )\n",
    "\n",
    "plt.imshow(cv_image)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23ee4e05",
   "metadata": {},
   "source": [
    "## Evaluation on the Test set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "3bdb32e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_true_object(root):\n",
    "    for child in root:\n",
    "        if child.tag == \"name\":\n",
    "            label_desc = str(child.text)\n",
    "            if label_desc == \"with_mask\":\n",
    "                return class_dict[\"with_mask\"]\n",
    "            elif label_desc == \"without_mask\":\n",
    "                return class_dict[\"without_mask\"]\n",
    "            else:\n",
    "                return class_dict[\"mask_weared_incorrect\"]\n",
    "\n",
    "def get_true_result(filename):\n",
    "    wi_m = 0\n",
    "    wo_m = 0\n",
    "    inc_m = 0\n",
    "    tree = ET.parse('./data/MaskedFace/val/' + filename + \".xml\")\n",
    "    root = tree.getroot()\n",
    "    result = None\n",
    "    for child in root:\n",
    "        if child.tag == \"object\":\n",
    "            result = read_true_object(child)\n",
    "            if result == class_dict[\"with_mask\"]:\n",
    "                wi_m+=1\n",
    "            elif result == class_dict[\"without_mask\"]:\n",
    "                wo_m+=1\n",
    "            elif result == class_dict[\"mask_weared_incorrect\"]:\n",
    "                inc_m+=1\n",
    "    return [wi_m, wo_m, inc_m]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "19b342c6",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████| 6/6 [05:57<00:00, 59.59s/it]\n"
     ]
    }
   ],
   "source": [
    "pred_list = []\n",
    "true_list = []\n",
    "\n",
    "model.eval()\n",
    "\n",
    "for images, targets, fname in tqdm(test_loader):\n",
    "    for i in range(len(fname)):\n",
    "        true_list.append(get_true_result(fname[i]))\n",
    "\n",
    "    images = list(image.to(device, dtype=torch.float) for image in images)\n",
    "    targets = [{k: v.to(device) for k, v in t.items()} for t in targets]\n",
    "\n",
    "    with torch.no_grad():\n",
    "        prediction = model(images)\n",
    "\n",
    "    for i in range(len(images)):\n",
    "        with_m = 0\n",
    "        with_o_m = 0\n",
    "        wrong_m = 0\n",
    "\n",
    "        for pred in range(len(prediction[i]['boxes'])):\n",
    "\n",
    "            box = prediction[i]['boxes'][pred].cpu().numpy()\n",
    "            label = prediction[i]['labels'][pred].cpu().numpy()\n",
    "            score = prediction[i]['scores'][pred].cpu().numpy()\n",
    "\n",
    "            if score >= 0.5:\n",
    "                if label.item() == class_dict[\"with_mask\"]:\n",
    "                    with_m += 1\n",
    "                elif label.item() == class_dict[\"without_mask\"]:\n",
    "                    with_o_m += 1\n",
    "                elif label.item() == class_dict[\"mask_weared_incorrect\"]:\n",
    "                    wrong_m += 1\n",
    "\n",
    "        pred_list.append([int(with_m), int(with_o_m), int(wrong_m)])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d77c192",
   "metadata": {},
   "source": [
    "## MAPE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "3734a3fb",
   "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": 15,
   "id": "dfa332ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "MAPE = compute_mape( np.array(pred_list), np.array(true_list) )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "8f6869dd",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "13.69502491681237\n"
     ]
    }
   ],
   "source": [
    "print(MAPE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83750954",
   "metadata": {},
   "source": [
    "## Final Score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "c6536ac1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Score:  21.875\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": "1409e328",
   "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
