{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Class\n",
    "\n",
    "0: head\n",
    "\n",
    "1: helmet\n",
    "\n",
    "2: face\n",
    "\n",
    "3: mask\n",
    "\n",
    "4: helmet & mask"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import cv2\n",
    "import math\n",
    "import numpy as np\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torchvision import transforms\n",
    "from sklearn.metrics import accuracy_score\n",
    "import yaml\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Classfication"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def crop_and_classify(img, boxes, classification_model, transform):\n",
    "    # classification_model.to(device)\n",
    "    # classification_model.eval()\n",
    "    \n",
    "    for box in boxes:\n",
    "        x1, y1, x2, y2 = box\n",
    "        # Expand the bounding box by 10%\n",
    "        width = x2 - x1\n",
    "        height = y2 - y1\n",
    "        x1 = max(int(x1 - 0.1 * width), 0)\n",
    "        y1 = max(int(y1 - 0.1 * height), 0)\n",
    "        x2 = min(int(x2 + 0.1 * width), img.shape[1] - 1)\n",
    "        y2 = min(int(y2 + 0.1 * height), img.shape[0] - 1)\n",
    "\n",
    "        # Crop the image\n",
    "        cropped_img = img[y1:y2, x1:x2]\n",
    "\n",
    "        # Transform the image to fit the input size of the classification model\n",
    "        transformed_img = transform(cropped_img)\n",
    "        transformed_img = transformed_img.unsqueeze(0)\n",
    "\n",
    "        # Classify the cropped image\n",
    "        with torch.no_grad():\n",
    "            outputs = classification_model(transformed_img)\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            confidence = torch.nn.functional.softmax(outputs, dim=1)[0][predicted].item()\n",
    "\n",
    "        # Print the classification result\n",
    "        print(f\"Class: {predicted.item()}, Confidence: {confidence:.2f}\")\n",
    "\n",
    "        # Draw the bounding box and classification result on the image\n",
    "        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)\n",
    "        cv2.putText(img, f\"Class: {predicted.item()}, Confidence: {confidence:.2f}\", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0,255,0), 2)\n",
    "    return img"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Detection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_detection_model(img, model_detect):\n",
    "    # img = transforms.ToTensor()(img).unsqueeze(0).to(device)  # Convert image to tensor and add batch dimension\n",
    "    # model_detect.eval()  # Set model to evaluation mode\n",
    "    img = transforms.ToTensor()(img).unsqueeze(0)\n",
    "    \n",
    "    # Move the model to the GPU\n",
    "    model_detect = model_detect.to(device)\n",
    "\n",
    "    # Run inference\n",
    "    with torch.no_grad():\n",
    "        detections = model_detect(img)\n",
    "\n",
    "    # Extract bounding boxes and confidence scores\n",
    "    boxes = detections.pred[0][:, :4].cpu().numpy().tolist()\n",
    "    scores = detections.pred[0][:, 4].cpu().numpy().tolist()\n",
    "\n",
    "    return boxes, scores\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load detection model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using cache found in C:\\Users\\pinb/.cache\\torch\\hub\\ultralytics_yolov5_master\n",
      "YOLOv5  2023-6-12 Python-3.10.9 torch-1.12.1 CUDA:0 (NVIDIA GeForce RTX 3060, 12288MiB)\n",
      "\n",
      "Fusing layers... \n",
      "Model summary: 157 layers, 7012822 parameters, 0 gradients, 15.8 GFLOPs\n",
      "Adding AutoShape... \n"
     ]
    }
   ],
   "source": [
    "weights_detect = \"C:/Users/pinb/Desktop/pytorch/yolov5/runs/train/dduk_64_/weights/best.pt\"\n",
    "model_detect = torch.hub.load('ultralytics/yolov5', 'custom', path=weights_detect)\n",
    "model_detect = model_detect.to(device).eval()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load classification model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using cache found in C:\\Users\\pinb/.cache\\torch\\hub\\ultralytics_yolov5_master\n",
      "YOLOv5  2023-6-12 Python-3.10.9 torch-1.12.1 CUDA:0 (NVIDIA GeForce RTX 3060, 12288MiB)\n",
      "\n",
      "Fusing layers... \n",
      "Model summary: 117 layers, 4173093 parameters, 0 gradients, 10.4 GFLOPs\n",
      "WARNING  YOLOv5 ClassificationModel is not yet AutoShape compatible. You must pass torch tensors in BCHW to this model, i.e. shape(1,3,224,224).\n"
     ]
    }
   ],
   "source": [
    "weights_cls = \"C:/Users/pinb/Desktop/pytorch/yolov5/runs/train-cls/dduk_cls3/weights/best.pt\"\n",
    "model_cls = torch.hub.load('ultralytics/yolov5', 'custom', path=weights_cls)\n",
    "model_cls = model_cls.to(device).eval()\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load YAML"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # YAML 파일 경로\n",
    "yaml_path = \"C:/Users/pinb/Desktop/pytorch/yolov5/runs/detect/test/data.yaml\"\n",
    "\n",
    "def read_dataset_paths(file_path):\n",
    "    with open(file_path, 'r') as file:\n",
    "        data = yaml.safe_load(file)\n",
    "    test_path = data.get('test')\n",
    "    train_path = data.get('train')\n",
    "    val_path = data.get('val')\n",
    " \n",
    "    return  os.path.join(os.path.dirname(file_path), val_path)\n",
    "\n",
    "# # YAML 파일에서 데이터셋 경로 읽기\n",
    "val_path = read_dataset_paths(yaml_path)\n",
    "\n",
    "# 이미지가 있는 폴더 경로\n",
    "folder_path = val_path\n",
    "result_path = os.path.join(folder_path, \"result\")\n",
    "# folder_path = \"C:/Users/pinb/Desktop/pytorch/yolov5/runs/detect/test/\"\n",
    "# result_path = \"C:/Users/pinb/Desktop/pytorch/yolov5/runs/detect/test/result\"\n",
    "\n",
    "def get_image_and_label_paths(folder_path):\n",
    "    image_paths = []\n",
    "    label_paths = []\n",
    "    for file_name in os.listdir(folder_path):\n",
    "        if file_name.endswith(\".bmp\") or file_name.endswith(\".jpg\") or file_name.endswith(\".png\"):\n",
    "            image_path = os.path.join(folder_path, file_name)\n",
    "            image_paths.append(image_path)\n",
    "            file_name = os.path.splitext(file_name)[0] + \".txt\"\n",
    "            label_path = folder_path.replace(\"images\", \"labels\")\n",
    "            label_path = os.path.join(label_path, file_name)\n",
    "            if os.path.exists(label_path):\n",
    "                label_paths.append(label_path)\n",
    "            else:\n",
    "                label_paths.append(\"\")\n",
    "\n",
    "    return image_paths, label_paths\n",
    "\n",
    "# 폴더 내의 이미지 파일 경로들을 가져옴\n",
    "image_paths, label_paths = get_image_and_label_paths(folder_path)\n",
    "\n",
    "if not os.path.exists(result_path):\n",
    "    os.makedirs(result_path)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Read Label Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_label(label_path):\n",
    "    labels = []\n",
    "    with open(label_path, 'r') as file:\n",
    "        lines = file.readlines()\n",
    "        for line in lines:\n",
    "            label = line.strip().split(' ')\n",
    "            label_class = int(label[0])\n",
    "            bounding_box = tuple(map(float, label[1:]))\n",
    "            labels.append((label_class, bounding_box))\n",
    "    return labels"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Find closet label class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_closest_label_class(pred_box, labels):\n",
    "    min_distance = float('inf')\n",
    "    closest_class = -1\n",
    "\n",
    "    for i, label in enumerate(labels):\n",
    "        label_box = label[1]\n",
    "        label_center_x = (label_box[0] + label_box[2]) / 2\n",
    "        label_center_y = (label_box[1] + label_box[3]) / 2\n",
    "\n",
    "        pred_center_x = (pred_box[0] + pred_box[2]) / 2\n",
    "        pred_center_y = (pred_box[1] + pred_box[3]) / 2\n",
    "\n",
    "        distance = math.sqrt((label_center_x - pred_center_x) ** 2 + (label_center_y - pred_center_y) ** 2)\n",
    "\n",
    "        if distance < min_distance:\n",
    "            min_distance = distance\n",
    "            closest_class = label[0]\n",
    "\n",
    "    return closest_class"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "lst_result = [[],[],[],[],[]] # class 0~4\n",
    "eval_times = []\n",
    "\n",
    "def Evaluate(model_detect, model_cls, img, img_name, labels, output=True):\n",
    "    height, width, _ = img.shape\n",
    "    expand_ratio = 0.5  # The ratio to expand the ROI area\n",
    "\n",
    "    # Inference detection model\n",
    "    start_time = time.time()\n",
    "    results = model_detect(img)\n",
    "\n",
    "    # For each detection, crop the ROI and classify it\n",
    "    for *box, detect_conf, cls in results.xyxy[0]:\n",
    "        x1, y1, x2, y2 = map(int, box)\n",
    "        fbox = map(float, box)\n",
    "\n",
    "        # Calculate the width and height of the bounding box\n",
    "        bbox_width = x2 - x1\n",
    "        bbox_height = y2 - y1\n",
    "        roi = (x1, y1, bbox_width, bbox_height)\n",
    "\n",
    "        # Calculate the expanded coordinates\n",
    "        x1_expanded = max(0, x1 - int(bbox_width * expand_ratio))\n",
    "        y1_expanded = max(0, y1 - int(bbox_height * expand_ratio))\n",
    "        x2_expanded = min(width, x2 + int(bbox_width * expand_ratio))\n",
    "        y2_expanded = min(height, y2 + int(bbox_height * expand_ratio))\n",
    "\n",
    "        roi_flat = (x1_expanded, \n",
    "                    y1_expanded, \n",
    "                    bbox_width + int(bbox_width * expand_ratio) * 2, \n",
    "                    bbox_height + int(bbox_height * expand_ratio) * 2)\n",
    "\n",
    "        # Crop the expanded ROI from the image\n",
    "        # roi = img[y1:y2, x1:x2]\n",
    "        roi_expanded = img.copy()\n",
    "        roi_expanded = roi_expanded[y1_expanded:y2_expanded, x1_expanded:x2_expanded]\n",
    "        roi_expanded = cv2.resize(roi_expanded, (224, 224))\n",
    "\n",
    "        # Convert numpy array to torch tensor, and normalize pixel values\n",
    "        roi_expanded = torch.from_numpy(roi_expanded).float().div(255.0)\n",
    "\n",
    "        # Reshape tensor to (channels, height, width)\n",
    "        roi_expanded = roi_expanded.permute(2, 0, 1)\n",
    "\n",
    "        # Add an extra dimension for the batch\n",
    "        roi_expanded = roi_expanded.unsqueeze(0)\n",
    "\n",
    "        # Move roi_expanded to the same device as class_model\n",
    "        roi_expanded = roi_expanded.to(device)\n",
    "        class_result = model_cls(roi_expanded)\n",
    "\n",
    "        # classfication result\n",
    "        probabilities = F.softmax(class_result, dim=1)\n",
    "        max_confidence, max_indices = torch.max(probabilities, dim=1)\n",
    "        class_pred = max_indices.item()\n",
    "        class_conf = max_confidence.item()\n",
    "\n",
    "        # confidence\n",
    "        total_conf = (detect_conf + class_conf) * 0.5\n",
    "\n",
    "        # label - class\n",
    "        class_gt = find_closest_label_class((x1, y1, x2, y2), labels)\n",
    "\n",
    "        # append (class, confidence, bounding box)\n",
    "        lst_result[class_pred].append((class_pred, class_gt, total_conf, (x1, y1, x2, y2)))\n",
    "\n",
    "        # Put classification result on each ROI\n",
    "        thin = 1\n",
    "        color = ()\n",
    "        if class_pred == 0: # head\n",
    "            color = (0, 0, 255)\n",
    "        elif class_pred == 1: # helmet\n",
    "            color = (255, 0, 0)\n",
    "        elif class_pred == 2: # face\n",
    "            color = (0, 255, 0)\n",
    "        elif class_pred == 3: # mask\n",
    "            color = (255, 0, 255)\n",
    "        elif class_pred == 4: # helmet & mask\n",
    "            color = (255, 255, 0)\n",
    "        cv2.putText(img, f'{class_pred}: {total_conf}', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, thin)\n",
    "        # cv2.rectangle(img, roi_flat, color, thin)\n",
    "        cv2.rectangle(img, roi, color, thin)\n",
    "        if output is True:\n",
    "            print(f'{img_name}\\'s det + cls result [{class_pred}]: {total_conf}%, {roi}:')\n",
    "\n",
    "    end_time = time.time()\n",
    "    total_time = end_time - start_time\n",
    "    eval_times.append(total_time)\n",
    "    if output is True:\n",
    "        print(f'Detect + Classfication Time: {total_time}s')\n",
    "\n",
    "    # Save image\n",
    "    cv2.imwrite(f'{result_path}/{img_name}_rst.jpg', img)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Compare Ground Truth & Bounding Box"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Class\n",
    "# 0: head\n",
    "# 1: helmet\n",
    "# 2: face\n",
    "# 3: mask\n",
    "# 4: helmet & mask\n",
    "\n",
    "def CompareClass(pr, gt):\n",
    "    if gt == 0: # head or face\n",
    "        if pr == 0 or pr == 2:\n",
    "            return True\n",
    "        else:\n",
    "            False\n",
    "    elif gt == 1: # helmet or mask\n",
    "        if pr == 1 or pr == 3 or pr == 4:\n",
    "            return True\n",
    "        else:\n",
    "            return False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [],
   "source": [
    "def CompareGTandBox(predict, setting_mAP=0):\n",
    "    # 0: mAP@.5, 1: mAP@.5:.95 0\n",
    "    # 비교하여 TP와 FP 구분\n",
    "    results = []\n",
    "\n",
    "    tp_cnt = 0\n",
    "    fp_cnt = 0\n",
    "    totalbox_cnt = 0\n",
    "    targetbox_cnt = len(predict)\n",
    "\n",
    "    for p_class, gt_class, confidence, p_box in predict:\n",
    "        if CompareClass(p_class, gt_class):\n",
    "            # iou = calculate_iou(p_box, l_box)\n",
    "            # iou = confidence\n",
    "            if setting_mAP == 0: # mAP@.5\n",
    "                if confidence >= 0.5:\n",
    "                    # True Positive (TP)\n",
    "                    tp_cnt = tp_cnt + 1\n",
    "                    totalbox_cnt = totalbox_cnt +1\n",
    "                    # class, confidence, boudingbox, TP/FP, precision, recall\n",
    "                    results.append((p_class, confidence, p_box, True, tp_cnt/totalbox_cnt, tp_cnt/targetbox_cnt))\n",
    "            elif setting_mAP == 1: # mAP@.5:.95\n",
    "                if confidence >= 0.5 and confidence < 0.95:\n",
    "                    # True Positive (TP)\n",
    "                    tp_cnt = tp_cnt + 1\n",
    "                    totalbox_cnt = totalbox_cnt + 1\n",
    "                    # class, confidence, boudingbox, TP/FP, precision, recall\n",
    "                    results.append((p_class, confidence, p_box, True, tp_cnt/totalbox_cnt, tp_cnt/targetbox_cnt))\n",
    "            else: # all P, R\n",
    "                tp_cnt = tp_cnt + 1\n",
    "                totalbox_cnt = totalbox_cnt + 1\n",
    "                results.append((p_class, confidence, p_box, True, tp_cnt/totalbox_cnt, tp_cnt/targetbox_cnt))\n",
    "            \n",
    "        else:\n",
    "            # False Positive (FP)\n",
    "            totalbox_cnt = totalbox_cnt + 1\n",
    "            results.append((p_class, confidence, p_box, False, tp_cnt/totalbox_cnt, tp_cnt/targetbox_cnt))\n",
    "\n",
    "    return results"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Get Precisions, Recalls and PR Curve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "def GetPR(results):\n",
    "    results = sorted(results, key=lambda x: x[5], reverse=True)\n",
    "    # results = sorted(results, key=lambda x: x[1], reverse=True)\n",
    "    precisions = [item[4] for item in results]\n",
    "    recalls = [item[5] for item in results]\n",
    "\n",
    "    return precisions, recalls"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Calculate AP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_AP(precisions, recalls):\n",
    "    assert len(recalls) == len(precisions)\n",
    "\n",
    "    sorted_indices = sorted(range(len(recalls)), key=lambda i: recalls[i])\n",
    "    sorted_precisions = [precisions[i] for i in sorted_indices]\n",
    "    sorted_recalls = [recalls[i] for i in sorted_indices]\n",
    "\n",
    "    ap = 0\n",
    "    for i in range(1, len(sorted_recalls)):\n",
    "        recall_diff = sorted_recalls[i] - sorted_recalls[i-1]\n",
    "        ap += recall_diff * (sorted_precisions[i] + sorted_precisions[i-1]) * 0.5\n",
    "\n",
    "    return ap"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Class 0\n",
      "meanP: 0.2451500387154128\n",
      "meanR: 0.1266693073789041\n",
      "AP@.5: 0.056300405455419165\n",
      "AP@.5:.95: 0.056300405455419165\n",
      "\n",
      "Class 1\n",
      "meanP: 0.8831870385742567\n",
      "meanR: 0.43243055555555554\n",
      "AP@.5: 0.6344911937169155\n",
      "AP@.5:.95: 0.6344911937169155\n",
      "\n",
      "Class 2\n",
      "meanP: 0\n",
      "meanR: 0\n",
      "AP@.5: 0\n",
      "AP@.5:.95: 0\n",
      "\n",
      "Class 3\n",
      "meanP: 0.8349450369473351\n",
      "meanR: 0.40800165622897366\n",
      "AP@.5: 0.591933273620593\n",
      "AP@.5:.95: 0.591933273620593\n",
      "\n",
      "Class 4\n",
      "meanP: 0.8888927177286716\n",
      "meanR: 0.43473958333333357\n",
      "AP@.5: 0.727062010345605\n",
      "AP@.5:.95: 0.727062010345605\n",
      "\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 49556 characters]",
      "text/plain": [
       "<Figure size 800x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 52636 characters]",
      "text/plain": [
       "<Figure size 800x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 29508 characters]",
      "text/plain": [
       "<Figure size 800x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 56212 characters]",
      "text/plain": [
       "<Figure size 800x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 58716 characters]",
      "text/plain": [
       "<Figure size 800x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==> Union Model (Detection + Classification) - 5 classes\n",
      "[mean Precisions] 0.5704349663931353\n",
      "[mean Recalls] 0.2803682204993534\n",
      "[mAP@.5] 0.40195737662770653\n",
      "[mAP@.5:.95] 0.40195737662770653\n",
      "[mean time] 0.042649137104018055s\n",
      "\n",
      "==> Convert to Model 1\n",
      "[mean Precisions] 0.6230357850195234\n",
      "[mean Recalls] 0.3059043619902019\n",
      "[mAP@.5] 0.4484272475282988\n",
      "[mAP@.5:.95] 0.4484272475282988\n",
      "[mean time] 0.042649137104018055s\n",
      "\n",
      "==> Convert to Model 2\n",
      "[mean Precisions] 0.6190156182172799\n",
      "[mean Recalls] 0.30386862037965345\n",
      "[mAP@.5] 0.3349644805230888\n",
      "[mAP@.5:.95] 0.8897615083738772\n",
      "[mean time] 0.042649137104018055s\n",
      "\n"
     ]
    }
   ],
   "source": [
    "lst_result = [[],[],[],[],[]] # class 0~4\n",
    "eval_times = []\n",
    "\n",
    "for i in range(0, len(image_paths)):\n",
    "    # Image name\n",
    "    img_dir = os.path.dirname(image_paths[i])\n",
    "    img_name = os.path.basename(image_paths[i])\n",
    "\n",
    "    # Load image\n",
    "    img = cv2.imread(image_paths[i]) \n",
    "    img = cv2.resize(img, (640, 640))\n",
    "    labels = read_label(label_paths[i])\n",
    "    Evaluate(model_detect, model_cls, img, img_name, labels, False)\n",
    "\n",
    "Ps = []\n",
    "Rs = []\n",
    "ap5s = []\n",
    "ap95s = []\n",
    "\n",
    "# Class\n",
    "# 0: head\n",
    "# 1: helmet\n",
    "# 2: face\n",
    "# 3: mask\n",
    "# 4: helmet & mask\n",
    "\n",
    "# each class result\n",
    "class_len = len(lst_result)\n",
    "\n",
    "for i in range(0, class_len): \n",
    "    results_all = CompareGTandBox(lst_result[i], -1) # all P, R\n",
    "    results_5 = CompareGTandBox(lst_result[i], 0) # mAP@.5\n",
    "    results_95 = CompareGTandBox(lst_result[i], 1) # mAP@.5:.95\n",
    "\n",
    "    p, r = GetPR(results_all)\n",
    "    p1, r1 = GetPR(results_5)\n",
    "    p2, r2 = GetPR(results_95)\n",
    "    ap1 = calculate_AP(p1, r1)\n",
    "    ap2 = calculate_AP(p2, r2)\n",
    "    # print(p1)\n",
    "    # print(r1)\n",
    "    # print(p2)\n",
    "    # print(r2)\n",
    "    if len(p) != 0:\n",
    "        p = sum(p) / len(p)\n",
    "    else:\n",
    "        p = 0\n",
    "    if len(r) != 0:\n",
    "        r = sum(r) / len(r)\n",
    "    else:\n",
    "        r = 0\n",
    "\n",
    "    Ps.append(p)\n",
    "    Rs.append(r)\n",
    "    ap5s.append(ap1)\n",
    "    ap95s.append(ap2)\n",
    "    print(f'Class {i}')\n",
    "    print(f'meanP:', p)\n",
    "    print(f'meanR:', r)\n",
    "    print(f\"AP@.5:\", ap1)\n",
    "    print(f\"AP@.5:.95:\", ap2)\n",
    "\n",
    "    fig, axs = plt.subplots(1, 2, figsize=(8, 4))\n",
    "    axs[0].plot(r1, p1)\n",
    "    axs[0].set_xlabel('Recall')\n",
    "    axs[0].set_ylabel('Precision')\n",
    "    axs[1].plot(r2, p2)\n",
    "    axs[1].set_xlabel('Recall')\n",
    "    axs[1].set_ylabel('Precision')\n",
    "    plt.suptitle(f'Class {i} Precision-Recall Curve (mAP@.5, mAP@.5:.95)')\n",
    "    print()\n",
    "    \n",
    "# plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "mP = sum(Ps) / class_len\n",
    "mR = sum(Rs) / class_len\n",
    "mAP_5 = sum(ap5s) / class_len\n",
    "mAP_95 = sum(ap95s) / class_len\n",
    "mean_time = sum(eval_times) / len(eval_times)\n",
    "\n",
    "print(f'==> Union Model (Detection + Classification) - 5 classes')\n",
    "print(f'[mean Precisions] {mP}')\n",
    "print(f'[mean Recalls] {mR}')\n",
    "print(f'[mAP@.5] {mAP_5}')\n",
    "print(f'[mAP@.5:.95] {mAP_95}')\n",
    "print(f'[mean time] {mean_time}s')\n",
    "print()\n",
    "\n",
    "\n",
    "#### Convert to evaluate ####\n",
    "\n",
    "# Class\n",
    "# 0: head\n",
    "# 1: helmet\n",
    "# 2: face\n",
    "# 3: mask\n",
    "# 4: helmet & mask\n",
    "\n",
    "# => Model 1: (0, 2, 3) - head, (1, 4) - helmet \n",
    "# => Model 2: (0, 1, 2) - face, (3, 4) - mask \n",
    "\n",
    "print(f'==> Convert to Model 1')\n",
    "print(f'[mean Precisions] {((Ps[0] + Ps[2] + Ps[3]) / 3 + (Ps[1] + Ps[4]) / 2) / 2}')\n",
    "print(f'[mean Recalls] {((Rs[0] + Rs[2] + Rs[3]) / 3 + (Rs[1] + Rs[4]) / 2) / 2}')\n",
    "print(f'[mAP@.5] {((ap5s[0] + ap5s[2] + ap5s[3]) / 3 + (ap5s[1] + ap5s[4]) / 2) / 2}')\n",
    "print(f'[mAP@.5:.95] {((ap95s[0] + ap95s[2] + ap95s[3]) / 3 + (ap95s[1] + ap95s[4]) / 2) / 2}')\n",
    "print(f'[mean time] {mean_time}s')\n",
    "print()\n",
    "\n",
    "print(f'==> Convert to Model 2')\n",
    "print(f'[mean Precisions] {((Ps[0] + Ps[1] + Ps[2]) / 3 + (Ps[3] + Ps[4]) / 2) / 2}')\n",
    "print(f'[mean Recalls] {((Rs[0] + Rs[1] + Rs[2]) / 3 + (Rs[3] + Rs[4]) / 2) / 2}')\n",
    "print(f'[mAP@.5] {((ap5s[0] + ap5s[1] + ap5s[2]) / 3 + (ap5s[3] + ap5s[4]) / 3) / 2}')\n",
    "print(f'[mAP@.5:.95] {((ap95s[0] + ap95s[1] + ap95s[2]) / 3 + (ap95s[3] + ap95s[4]) / 2)}')\n",
    "print(f'[mean time] {mean_time}s')\n",
    "print()\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
