IonoDet.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mmcv 2.0.0rc4\n",
      "mmengine 0.6.0\n",
      "mmdet 3.0.0rc6\n",
      "mmseg 1.0.0rc6\n",
      "mmyolo 0.5.0\n"
     ]
    }
   ],
   "source": [
    "import mmcv\n",
    "import mmengine\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "from mmdet.apis import init_detector, inference_detector\n",
    "from tabulate import tabulate\n",
    "import matplotlib.patches as mpatches\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "import cv2\n",
    "from mmseg.apis import inference_model, init_model\n",
    "import mmdet\n",
    "import mmseg\n",
    "import mmyolo\n",
    "import copy\n",
    "print('mmcv', mmcv.__version__)\n",
    "print('mmengine', mmengine.__version__)\n",
    "print('mmdet', mmdet.__version__)\n",
    "print('mmseg', mmseg.__version__)\n",
    "print('mmyolo', mmyolo.__version__)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model Initialization"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Init Segmentor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/mmsegmentation/mmseg/models/builder.py:36: UserWarning: ``build_loss`` would be deprecated soon, please use ``mmseg.registry.MODELS.build()`` \n",
      "  warnings.warn('``build_loss`` would be deprecated soon, please use '\n",
      "/home/ubuntu/mmsegmentation/mmseg/models/losses/cross_entropy_loss.py:235: UserWarning: Default ``avg_non_ignore`` is False, if you would like to ignore the certain label and average loss over non-ignore labels, which is the same with PyTorch official cross_entropy, set ``avg_non_ignore=True``.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loads checkpoint by local backend from path: ../mmsegmentation/work_dirs/se4ionogram/pspnet_r50_ionogram_iou_3922_acc_9153.pth\n"
     ]
    }
   ],
   "source": [
    "data_root = '../mmsegmentation/data/IonoSeg/'\n",
    "img_dir = 'rgbimg'\n",
    "ann_dir = 'rgbmask'\n",
    "classes = ('Background', 'E', 'Es-l', 'Es-c', 'F1', 'F2', 'Spread-F')\n",
    "palette = [[255, 255, 255], [250, 165, 30], [120, 69, 125], [53, 125, 34], \n",
    "           [0, 11, 123], [130, 20, 12], [120, 121, 80]] # F1 [78, 210, 240]\n",
    "patches = [mpatches.Patch(color=np.array(palette[i]) / 255., label=classes[i]) for i in range(7)]\n",
    "\n",
    "seg_config_file = '../mmsegmentation/configs/_se4ionogram/pspnet_r50_ionogram_mmseg1.py'\n",
    "seg_checkpoint_file = '../mmsegmentation/work_dirs/se4ionogram/pspnet_r50_ionogram_iou_3922_acc_9153.pth'\n",
    "\n",
    "seg_model = init_model(seg_config_file, seg_checkpoint_file, device='cuda:0')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mmseg.registry import DATASETS\n",
    "from mmseg.datasets import BaseSegDataset\n",
    "\n",
    "\n",
    "@DATASETS.register_module()\n",
    "class IonogramSegmentationDataset(BaseSegDataset):\n",
    "  METAINFO = dict(classes = classes, palette = palette)\n",
    "  def __init__(self, **kwargs):\n",
    "    super().__init__(img_suffix='.png', seg_map_suffix='.png', **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def single_load(model, img):\n",
    "    segs = inference_model(model, img)\n",
    "    segmap = segs.pred_sem_seg.data.cpu().squeeze(0).numpy()\n",
    "    return segmap\n",
    "\n",
    "def post_seg(segmap, img):\n",
    "    blurred = cv2.GaussianBlur(img, (3, 3), 0)\n",
    "    edge = cv2.Canny(blurred, 50, 150)\n",
    "    edge = edge / 255\n",
    "    segmap = segmap * edge   # Hadamard product\n",
    "    return segmap\n",
    "\n",
    "def interpret_numpy(segmap):\n",
    "    labels = []\n",
    "    parameters = []\n",
    "    for class_id in range(1, 7):\n",
    "        indices = np.where(segmap == class_id)\n",
    "        if len(indices[0]) != 0:\n",
    "            parameters.append([np.max(indices[1]), np.max(indices[0])])\n",
    "            labels.append(class_id - 1)\n",
    "    # # fmin\n",
    "    # indices = np.nonzero(segmap)\n",
    "    # if len(indices[0]) != 0:\n",
    "    #     fmin = np.min(indices[1])\n",
    "    return [torch.tensor(parameters), labels]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Init Detector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loads checkpoint by local backend from path: ./work_dirs/rtmdet_tiny_100e/best_coco/bbox_mAP_epoch_100_0.589.pth\n",
      "{'multi_label': True, 'nms_pre': 30000, 'score_thr': 0.001, 'nms': {'type': 'nms', 'iou_threshold': 0.65}, 'max_per_img': 100}\n",
      "646 images in the test set\n"
     ]
    }
   ],
   "source": [
    "val_dir = './Iono4311/val_images'\n",
    "test_dir = './Iono4311/test_images'\n",
    "ann_val = './Iono4311/annotations/val.json'\n",
    "ann_test = './Iono4311/annotations/test.json'\n",
    "\n",
    "# Choose to use a config and initialize the detector\n",
    "det_config_file = './configs/custom_dataset/rtmdet/rtmdet_tiny_fast_1xb32-100e_ionogram.py'\n",
    "# Setup a checkpoint file to load\n",
    "det_checkpoint_file = './work_dirs/rtmdet_tiny_100e/best_coco/bbox_mAP_epoch_100_0.589.pth'\n",
    "\n",
    "# build the model from a config file and a checkpoint file\n",
    "det_model = init_detector(det_config_file, det_checkpoint_file)\n",
    "det_model.cfg['model']['test_cfg']['nms']['iou_threshold'] = 0.65\n",
    "print(det_model.cfg['model']['test_cfg'])\n",
    "\n",
    "annotation = mmengine.load(ann_test)\n",
    "name_list = []\n",
    "for image in annotation['images']:\n",
    "    name_list.append(image['file_name'][:-4])\n",
    "print(f'{len(name_list)} images in the test set')\n",
    "# 建立 name2index 索引\n",
    "name2annid = {}\n",
    "for image in annotation['images']:\n",
    "    ann_id = []\n",
    "    for i, ann in enumerate(annotation['annotations']):\n",
    "        if ann['image_id'] == image['id']:\n",
    "            ann_id.append(i)\n",
    "    name2annid[image['file_name']] = ann_id\n",
    "\n",
    "img_names = []\n",
    "for image in annotation['images']:\n",
    "    img_names.append(image['file_name'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_detection_pred(model, img):\n",
    "    pred = inference_detector(model, img)\n",
    "    pred_instances = pred.pred_instances\n",
    "    pred_instances = pred_instances[pred_instances.scores > 0.3]\n",
    "\n",
    "    # 每一类只保留一个置信度最高的 bbox\n",
    "    max_idx = []\n",
    "    for i in range(torch.max(pred_instances.labels) + 1):\n",
    "        idx = torch.where(pred_instances.labels == i)[0]\n",
    "        if len(idx) > 0:\n",
    "            max_idx.append(idx[torch.argmax(pred_instances.scores[idx])])\n",
    "    max_idx = torch.tensor(max_idx)\n",
    "    scores = pred_instances.scores[max_idx].tolist()\n",
    "    labels = pred_instances.labels[max_idx].tolist()\n",
    "    bboxes = pred_instances.bboxes[max_idx].cpu()\n",
    "    params = bboxes[:, 2:4] # [height, frequency]\n",
    "    return bboxes, params, labels, scores\n",
    "\n",
    "def get_segmentation_pred(model, img):\n",
    "    segmap = single_load(model, img)\n",
    "    segmap = post_seg(segmap, img)\n",
    "    parameters, labels = interpret_numpy(segmap)\n",
    "    return [segmap, parameters, labels]\n",
    "\n",
    "def interprete_parameters(boundaries, labels):\n",
    "    # [hE, foE, hEs, foEs, hF1, foF1, hF2, foF2] = [-1] * 8\n",
    "    params = [-1] * 8\n",
    "    # labels: parameter_id      labels begin with 0\n",
    "    boundary2param = {\n",
    "        0: 0,       # E\n",
    "        1: 2, 2: 2, # Es\n",
    "        3: 4,       # F1\n",
    "        4: 6, 5: 6  # F2, Spread-F\n",
    "    }\n",
    "    for i in range(len(labels)):\n",
    "        params[boundary2param[labels[i]]: boundary2param[labels[i]] + 2] = boundaries[i].tolist()\n",
    "    return params\n",
    "\n",
    "def get_gt(file_name):\n",
    "    # bbox [x1, y1, x2, y2] \n",
    "    # x1 refers to the distance between the upperleft point and left boundary,\n",
    "    # y1 refers to the distance between the upperleft point and the top boundary.\n",
    "    labels = [annotation['annotations'][i]['category_id'] - 1 for i in name2annid[file_name]]   # labels begin with 0\n",
    "    bboxes = [annotation['annotations'][i]['bbox'] for i in name2annid[file_name]]\n",
    "    bboxes = torch.tensor(bboxes)\n",
    "    parameters = [[box[0] + box[2], box[1] + box[3]] for box in bboxes]\n",
    "    parameters = torch.tensor(parameters)\n",
    "    bboxes[:, 2:4] = parameters\n",
    "    mask = Image.open(f'{data_root}/{ann_dir}/{file_name}')\n",
    "    return [mask, bboxes, parameters, labels]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "colors = [[0.98, 0.647, 0.118], [0.471, 0.271, 0.49], [0.208, 0.49, 0.133], [0.307, 0.823, 0.941], [0.51, 0.078, 0.047], [0.471, 0.475, 0.314]]\n",
    "def show_bboxes(axes, bboxes, catogories, labels=None):\n",
    "\n",
    "    def make_list(obj, default_values=None):\n",
    "        if obj is None:\n",
    "            obj = default_values\n",
    "        elif not isinstance(obj, (list, tuple)):\n",
    "            obj = [obj]\n",
    "        return obj\n",
    "\n",
    "    labels = make_list(labels)\n",
    "    for i, bbox in enumerate(bboxes):\n",
    "        color = colors[catogories[i]]\n",
    "        rect = plt.Rectangle(\n",
    "            xy=(bbox[0], bbox[1]), width=bbox[2]-bbox[0],\n",
    "            height=bbox[3]-bbox[1], fill=False, edgecolor=color, linewidth=1.5)\n",
    "        axes.add_patch(rect)\n",
    "        if labels and len(labels) > i:\n",
    "            text_color = 'w'\n",
    "            axes.text(rect.xy[0], rect.xy[1] - 6, labels[i],\n",
    "                      va='center', ha='center', fontsize=9, color=text_color,\n",
    "                      bbox=dict(facecolor=color, lw=0, pad=0.5))\n",
    "\n",
    "def set_ticks(ax):\n",
    "    ax.set_xticks([0, 50, 100, 150, 200, 250, 300, 350, 400], [1.0, 3.5, 6.0, 8.5, 11.0, 13.5, 16.0, 18.5, 21.0])\n",
    "    ax.set_yticks([0, 40, 80, 120, 160, 200, 240, 280, 320, 360], [990, 890, 790, 690, 590, 490, 390, 290, 190, 90])\n",
    "    ax.set_xlabel('Frequency (MHz)')\n",
    "    ax.set_ylabel('Virtual Height (km)')\n",
    "\n",
    "def plot_layer_parameters(ax, layer_name, h, f, color='r', h_offset=0, f_offset=0):\n",
    "    if h > 0:\n",
    "        ax.axvline(f, color=color, linestyle=\":\", alpha=0.8)    # foX\n",
    "        ax.text(h_offset - 10, -3, f'fo{layer_name}', color=color)\n",
    "        ax.axhline(h, color=color, linestyle=\"--\", alpha=0.8)   # h'X\n",
    "        ax.text(401, f_offset + 2, f'h\\'{layer_name}', color=color)\n",
    "\n",
    "# params = [hE, foE, hEs, foEs, hF1, foF1, hF2, foF2]\n",
    "boundary2param = {\n",
    "    'E': {'offset': -6, 'param_index':[1, 0], 'color': [colors[0], colors[0]]},\n",
    "    'Es': {'offset': 0, 'param_index':[3, 2], 'color': [colors[2], colors[1]]},\n",
    "    'F1': {'offset': 0, 'param_index':[5, 4], 'color': [colors[3], colors[3]]},\n",
    "    'F2': {'offset': 0, 'param_index':[7, 6], 'color': [colors[4], colors[5]]}\n",
    "}\n",
    "\n",
    "def params2txtpos(params, width):\n",
    "    sorted_params = sorted(params)   # 递增\n",
    "    text_pos = copy.deepcopy(sorted_params)\n",
    "    l = len(params)\n",
    "    deviation = [sorted_params[i + 1] - sorted_params[i] for i in range(0, l-1)]\n",
    "    for i in range(len(deviation)):\n",
    "        if deviation[i] < width and deviation[i] != 0:\n",
    "            text_pos[i] -= (width-deviation[i]) / 2 \n",
    "            text_pos[i+1] += (width-deviation[i]) / 2 \n",
    "    mapping = dict(zip(sorted_params, text_pos))\n",
    "    resorted_lst = [mapping[x] for x in params]\n",
    "\n",
    "    return resorted_lst\n",
    "\n",
    "def plot_parameters(ax, params, labels):\n",
    "    txtpos_f = params2txtpos(params[1:8:2], width=14)\n",
    "    txtpos_h = params2txtpos(params[0:8:2], width=26)\n",
    "    for key in boundary2param.keys():\n",
    "        if key == 'Es' and 2 in labels:\n",
    "            color = boundary2param[key]['color'][0]\n",
    "        elif key == 'F2' and 4 in labels: # Es-c or F2\n",
    "            color = boundary2param[key]['color'][0]\n",
    "        else:\n",
    "            color = boundary2param[key]['color'][1]\n",
    "        plot_layer_parameters(ax, key, params[boundary2param[key]['param_index'][0]], params[boundary2param[key]['param_index'][1]], color=color, \n",
    "                              h_offset=txtpos_h[boundary2param[key]['param_index'][1] // 2], f_offset=txtpos_f[boundary2param[key]['param_index'][1] // 2])\n",
    "\n",
    "def show_det_result(file_name, bboxes, boundaries_pred, labels, figsize=(18, 6)):\n",
    "    _, axs = plt.subplots(1, 3, figsize=figsize)\n",
    "    img = mmcv.imread(test_dir + '/' + file_name)\n",
    "    [axs[i].imshow(255 - img) for i in range(3)]\n",
    "    axs[0].set_title('Input Ionogram', y=1.018)\n",
    "    axs[1].set_title('Pred Bboxes', y=1.018)\n",
    "    axs[2].set_title('Ground Truth', y=1.018)\n",
    "\n",
    "    if isinstance(bboxes, torch.Tensor):    # x1, y1, x2, y2\n",
    "        bboxes = bboxes.cpu()\n",
    "    mask, bboxes_gt, boundaries_gt, labels_gt = get_gt(file_name)\n",
    "    params_pred = interprete_parameters(boundaries_pred, labels)\n",
    "    params_gt = interprete_parameters(boundaries_gt, labels_gt) # [hE, foE, hEs, foEs, hF1, foF1, hF2, foF2]\n",
    "    for ax in axs:\n",
    "        set_ticks(ax)\n",
    "    plot_parameters(axs[1], params_pred, labels)\n",
    "    plot_parameters(axs[2], params_gt, labels_gt)\n",
    "    show_bboxes(axs[1], bboxes, catogories=labels, labels=[det_model.dataset_meta['classes'][labels[i]] for i in range(len(labels))])\n",
    "    show_bboxes(axs[2], bboxes_gt, catogories=labels_gt, labels=[det_model.dataset_meta['classes'][labels_gt[i]] for i in range(len(labels_gt))])\n",
    "\n",
    "def show_seg_result(file_name, segmap, boundaries_pred, labels, figsize=(18, 6)):\n",
    "    _, axs = plt.subplots(1, 3, figsize=figsize)\n",
    "    img = mmcv.imread(test_dir + '/' + file_name)\n",
    "    mask, _, boundaries_gt, labels_gt = get_gt(file_name)\n",
    "    axs[0].imshow(img)\n",
    "    axs[2].imshow(mask)\n",
    "    seg_res = Image.fromarray(np.uint8(segmap)).convert('P')\n",
    "    seg_res.putpalette(np.array(palette, dtype=np.uint8))\n",
    "    axs[1].imshow(np.array(seg_res.convert('RGB')))\n",
    "    axs[0].set_title('Input Ionogram', y=1.018)\n",
    "    axs[1].set_title('Seg Map', y=1.018)\n",
    "    axs[2].set_title('Ground Truth', y=1.018)\n",
    "    for ax in axs: \n",
    "        set_ticks(ax)\n",
    "    params_pred = interprete_parameters(boundaries_pred, labels)\n",
    "    params_gt = interprete_parameters(boundaries_gt, labels_gt) # [hE, foE, hEs, foEs, hF1, foF1, hF2, foF2]\n",
    "    plot_parameters(axs[1], params_pred, labels)\n",
    "    plot_parameters(axs[2], params_gt, labels_gt)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Custom Evaluate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def boundaries_within_total_threshold(boundaries_pred, boundaries_gt, threshold=10):\n",
    "    return torch.sum(torch.abs(boundaries_pred - boundaries_gt)) <= threshold\n",
    "\n",
    "def all_boundaries_within_threshold(boundaries_pred, boundaries_gt, threshold=4):\n",
    "    return torch.all(torch.abs(boundaries_pred - boundaries_gt)<= threshold)\n",
    "\n",
    "def calculate_metrics(TP, FP, num_gts):\n",
    "    precisions, recalls = [], []\n",
    "    for i in range(6):\n",
    "        if num_gts[i] == 0:\n",
    "            precisions.append(-1)\n",
    "            recalls.append(-1)\n",
    "        else:\n",
    "            precisions.append(TP[i] / (TP[i] + FP[i]))\n",
    "            recalls.append(TP[i] / num_gts[i])\n",
    "\n",
    "    # Calculate Metrics\n",
    "    mean_precision = sum(precisions) / 6\n",
    "    mean_recall = sum(recalls) / 6\n",
    "    F1_score = 2 * mean_precision * mean_recall / (mean_precision + mean_recall)\n",
    "    overall_precision = sum(TP) / (sum(TP) + sum(FP))\n",
    "    overall_recall = sum(TP) / sum(num_gts)\n",
    "    Layer_names = list(det_model.dataset_meta['classes'])\n",
    "    Layer_names.extend(['Mean', 'Overall'])\n",
    "    precisions.extend([mean_precision, overall_precision])\n",
    "    recalls.extend([mean_recall, overall_recall])\n",
    "    formatted_precisions = ['{:.4f}'.format(precision) for precision in precisions]\n",
    "    formatted_recalls = ['{:.4f}'.format(recall) for recall in recalls]\n",
    "    result_table = list(zip(Layer_names, formatted_precisions, formatted_recalls))\n",
    "    print(tabulate(result_table, headers=['Layer', 'Precision', 'Recall']))\n",
    "    print(f'F1_score            {F1_score:.4f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loads checkpoint by local backend from path: ./work_dirs/rtmdet_tiny_100e/best_coco/bbox_mAP_epoch_100_0.589.pth\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/.conda/envs/mmlab/lib/python3.8/site-packages/torch/functional.py:478: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at  /opt/conda/conda-bld/pytorch_1659484810403/work/aten/src/ATen/native/TensorShape.cpp:2894.)\n",
      "  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Layer       Precision    Recall\n",
      "--------  -----------  --------\n",
      "E              0.8845    0.9147\n",
      "Es-l           0.7553    0.5772\n",
      "Es-c           0.791     0.8346\n",
      "F1             0.8867    0.9226\n",
      "F2             0.9634    0.9711\n",
      "Spread-F       0.6176    0.9545\n",
      "Mean           0.8164    0.8625\n",
      "Overall        0.8955    0.9057\n",
      "F1_score            0.8388\n"
     ]
    }
   ],
   "source": [
    "TP, FP = [0] * 6, [0] * 6\n",
    "num_gts = [0] * 6\n",
    "cnt = 0\n",
    "sample_interval = 1\n",
    "det_model = init_detector(det_config_file, det_checkpoint_file)\n",
    "# seg_model = init_model(seg_config_file, seg_checkpoint_file, device='cuda:0')\n",
    "\n",
    "# for image in annotation['images']:\n",
    "while cnt < len(annotation['images']):\n",
    "    img = mmcv.imread(test_dir + '/' + img_names[cnt])\n",
    "    # Get detection predictions\n",
    "    bboxes, boundaries_pred, labels, scores = get_detection_pred(det_model, img)\n",
    "    # Get segmentation predictions\n",
    "    # segmap, boundaries_pred, labels = get_segmentation_pred(seg_model, img)\n",
    "\n",
    "    # Get the ground truth in lists\n",
    "    mask, bboxes_gt, boundaries_gt, labels_gt = get_gt(img_names[cnt])\n",
    "    \n",
    "    # Interprete Prameters\n",
    "    params_pred = interprete_parameters(boundaries_pred, labels)\n",
    "    params_gt = interprete_parameters(boundaries_gt, labels_gt)\n",
    "    # print(cnt)\n",
    "    # show_det_result(img_names[cnt], bboxes, boundaries_pred, labels)\n",
    "    # show_seg_result(img_names[cnt], segmap, boundaries_pred, labels)\n",
    "    \n",
    "    for label in labels_gt:\n",
    "        num_gts[label] += 1\n",
    "    for i in range(len(labels)):\n",
    "        if labels[i] in labels_gt and all_boundaries_within_threshold(boundaries_pred[i], boundaries_gt[labels_gt.index(labels[i])]):\n",
    "            TP[labels[i]] += 1\n",
    "        else:\n",
    "            FP[labels[i]] += 1\n",
    "    cnt += sample_interval\n",
    "    if cnt > len(annotation['images']):\n",
    "        break\n",
    "\n",
    "calculate_metrics(TP, FP, num_gts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\n0: E, Es-c, F1, F2\\n70: Es-l, F2\\n76: Spread-F\\n\\n先验知识:\\n单独存在的 Es-c -> Es-l\\nSpread-F + F2 -> 取置信度高的\\nE + Es-l -> E + Es-c \\n'"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "'''\n",
    "0: E, Es-c, F1, F2\n",
    "70: Es-l, F2\n",
    "76: Spread-F\n",
    "\n",
    "先验知识:\n",
    "单独存在的 Es-c -> Es-l\n",
    "Spread-F + F2 -> 取置信度高的\n",
    "E + Es-l -> E + Es-c \n",
    "'''"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**seg** \n",
    "Layer       Precision    Recall\n",
    "--------  -----------  --------\n",
    "E              0.7621    0.8089\n",
    "Es-l           0.4805    0.6016\n",
    "Es-c           0.5924    0.7323\n",
    "F1             0.7941    0.8182\n",
    "F2             0.9665    0.9711\n",
    "Spread-F       0.6522    0.6818\n",
    "Mean           0.7079    0.769\n",
    "Overall        0.8034    0.8532\n",
    "F1_score            0.7372\n",
    "\n",
    "**det**\n",
    "Layer       Precision    Recall\n",
    "--------  -----------  --------\n",
    "E              0.8845    0.9147\n",
    "Es-l           0.7553    0.5772\n",
    "Es-c           0.791     0.8346\n",
    "F1             0.8867    0.9226\n",
    "F2             0.9634    0.9711\n",
    "Spread-F       0.6176    0.9545\n",
    "Mean           0.8164    0.8625\n",
    "Overall        0.8955    0.9057\n",
    "F1_score            0.8388"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### COCO Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# results_COCO = []\n",
    "# for image in annotation['images']:\n",
    "#     img = mmcv.imread(test_dir + '/' + image['file_name'])\n",
    "#     pred = inference_detector(model, img)\n",
    "#     pred_instances = pred.pred_instances\n",
    "#     pred_instances = pred_instances[pred_instances.scores > 0.1]\n",
    "#     bboxs = pred_instances.bboxes\n",
    "#     bboxs[:, 2], bboxs[:, 3] =  bboxs[:, 2] - bboxs[:, 0], bboxs[:, 3] - bboxs[:, 1]\n",
    "#     bboxs = bboxs.tolist()\n",
    "#     scores = pred_instances.scores.tolist()\n",
    "#     labels = pred_instances.labels.tolist()\n",
    "#     for i in range(len(labels)):\n",
    "#         pred_dict = {\n",
    "#             \"image_id\": image['id'],\n",
    "#             \"category_id\": labels[i] + 1,\n",
    "#             \"bbox\": bboxs[i],   # [x, y, width, height]\n",
    "#             \"score\": scores[i]\n",
    "#         }\n",
    "#         results_COCO.append(pred_dict)\n",
    "# mmengine.dump(results_COCO, './results_COCO.json')\n",
    "# len(results_COCO)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from pycocotools.coco import COCO\n",
    "# from pycocotools.cocoeval import COCOeval\n",
    "\n",
    "# # 初始化 COCO API\n",
    "# annFile = ann_test  # COCO 数据集的标注文件路径\n",
    "# cocoGt = COCO(annFile)\n",
    "\n",
    "# # 载入检测结果\n",
    "# # resFile = './results_COCO.json'  # 检测结果文件路径\n",
    "# cocoDt = cocoGt.loadRes(results_COCO)\n",
    "\n",
    "# # 创建评估器\n",
    "# cocoEval = COCOeval(cocoGt, cocoDt, 'bbox')\n",
    "\n",
    "# # 运行评估并获取结果\n",
    "# cocoEval.evaluate()\n",
    "# cocoEval.accumulate()\n",
    "# cocoEval.summarize()\n",
    "\n",
    "# # 输出 precisions 和 recall\n",
    "# print('precisions: ', cocoEval.stats[0])\n",
    "# print('recall: ', cocoEval.stats[8])\n",
    "\n",
    "# for catId in cocoGt.getCatIds():\n",
    "#     coco_eval = COCOeval(cocoGt, cocoDt, 'bbox')\n",
    "#     coco_eval.params.catIds = [catId]\n",
    "#     coco_eval.evaluate()\n",
    "#     coco_eval.accumulate()\n",
    "#     coco_eval.summarize()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Draft"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from mmdet.registry import VISUALIZERS\n",
    "# visualizer = VISUALIZERS.build(det_model.cfg.visualizer)\n",
    "# visualizer.dataset_meta = det_model.dataset_meta\n",
    "\n",
    "# image = mmcv.imread(test_dir + '/' +  annotation['images'][0]['file_name'])\n",
    "# print('id ', annotation['images'][0]['id'], annotation['images'][0]['file_name'])\n",
    "# result = inference_detector(det_model, image)\n",
    "# print(result)\n",
    "\n",
    "# # show the results\n",
    "# visualizer.add_datasample(\n",
    "#     'result',\n",
    "#     image,\n",
    "#     data_sample=result,\n",
    "#     draw_gt = None,\n",
    "#     wait_time=0,\n",
    "# )\n",
    "# visualizer.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AP50: 0.4091\n",
      "Recall: 0.6000\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "def calculate_iou(box1, box2):\n",
    "    # 计算两个矩形框的交集和并集\n",
    "    x1 = max(box1[0], box2[0])\n",
    "    y1 = max(box1[1], box2[1])\n",
    "    x2 = min(box1[2], box2[2])\n",
    "    y2 = min(box1[3], box2[3])\n",
    "    intersection = max(0, x2 - x1) * max(0, y2 - y1)\n",
    "    union = (box1[2] - box1[0]) * (box1[3] - box1[1]) + (box2[2] - box2[0]) * (box2[3] - box2[1]) - intersection\n",
    "    \n",
    "    # 计算 IoU 值\n",
    "    iou = intersection / union if union > 0 else 0\n",
    "    return iou\n",
    "\n",
    "detections = np.array([\n",
    "    [50, 50, 100, 100, 0.9],\n",
    "    [150, 150, 200, 200, 0.8],\n",
    "    [75, 75, 125, 125, 0.7],\n",
    "    [130, 130, 180, 180, 0.6],\n",
    "    [90, 90, 140, 140, 0.5],\n",
    "    [10, 10, 60, 60, 0.4],\n",
    "    [200, 200, 250, 250, 0.3],\n",
    "    [170, 170, 220, 220, 0.2],\n",
    "    [110, 110, 160, 160, 0.1],\n",
    "    [20, 20, 70, 70, 0.0]\n",
    "])\n",
    "\n",
    "gt_boxes = np.array([\n",
    "    [60, 60, 110, 110],\n",
    "    [155, 155, 205, 205],\n",
    "    [80, 80, 130, 130],\n",
    "    [130, 130, 180, 180],\n",
    "    [100, 100, 150, 150]\n",
    "])\n",
    "\n",
    "# 按照置信度从高到低排序\n",
    "sorted_indices = np.argsort(-detections[:, 4])\n",
    "sorted_detections = detections[sorted_indices]\n",
    "\n",
    "# 计算每个检测结果与所有 ground truth 的 IoU 值\n",
    "ious = np.zeros((len(sorted_detections), len(gt_boxes)))\n",
    "for i, detection in enumerate(sorted_detections):\n",
    "    for j, gt_box in enumerate(gt_boxes):\n",
    "        iou = calculate_iou(detection[:4], gt_box)\n",
    "        ious[i, j] = iou\n",
    "\n",
    "# 对每个检测结果，选择与其 IoU 值最大的 ground truth 进行匹配\n",
    "max_ious = ious.max(axis=1)\n",
    "matched_indices = ious.argmax(axis=1)\n",
    "\n",
    "# 根据 IoU 阈值选择 TP 和 FP\n",
    "iou_threshold = 0.5\n",
    "tp_flags = max_ious >= iou_threshold\n",
    "fp_flags = max_ious < iou_threshold\n",
    "fp_flags |= matched_indices == -1\n",
    "\n",
    "# 将 TP 和 FP 的标记添加到检测结果数组中\n",
    "sorted_detections = np.hstack((sorted_detections, tp_flags[:, np.newaxis], fp_flags[:, np.newaxis]))\n",
    "\n",
    "# 计算 Precision 和 Recall\n",
    "tp_cumsum = np.cumsum(tp_flags)\n",
    "fp_cumsum = np.cumsum(fp_flags)\n",
    "recall = tp_cumsum / len(gt_boxes)\n",
    "precision = tp_cumsum / (tp_cumsum + fp_cumsum)\n",
    "\n",
    "# 计算 AP50 值\n",
    "ap50 = 0.\n",
    "for t in np.arange(0., 1.1, 0.1):\n",
    "    mask = recall >= t\n",
    "    if mask.any():\n",
    "        ap50 += np.max(precision[mask]) / 11.\n",
    "\n",
    "# 输出结果\n",
    "print('AP50: {:.4f}'.format(ap50))\n",
    "print('Recall: {:.4f}'.format(recall[-1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "#  Average Precision  (precisions) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.588\n",
    "#  Average Precision  (precisions) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.874\n",
    "#  Average Precision  (precisions) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.618\n",
    "#  Average Precision  (precisions) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.543\n",
    "#  Average Precision  (precisions) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.715\n",
    "#  Average Precision  (precisions) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.894\n",
    "#  Average Recall     (recalls) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.632\n",
    "#  Average Recall     (recalls) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.670\n",
    "#  Average Recall     (recalls) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.670\n",
    "#  Average Recall     (recalls) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.637\n",
    "#  Average Recall     (recalls) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.760\n",
    "#  Average Recall     (recalls) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.925\n",
    "\n",
    "# bbox_mAP_copypaste: 0.589 0.876 0.621 0.544 0.723 0.897"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mmlab",
   "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.8.16"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "3f0cac9b4c15a115bcc92a822a73b5ddff4c7104025f7602d5a374ddd4c7361d"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}