GGO Localization with Bounding Boxes and YOLOv5.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 2,
  "metadata": {
    "colab": {
      "name": "Bounding_Boxes.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "c648pGhlmw-z",
        "XdGluD13p4De"
      ]
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3.8.11 64-bit ('MPF': conda)"
    },
    "accelerator": "GPU",
    "language_info": {
      "name": "python",
      "version": "3.8.11"
    },
    "interpreter": {
      "hash": "3dd217612e3efbc53bda62975acccf69dc94b7f474970d20d9b0bbb03e9d7037"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "from google.colab import drive\r\n",
        "drive.mount('/content/drive')"
      ],
      "outputs": [],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_OIxfe5Yi5Ag",
        "outputId": "b3f5aba9-818f-431b-cd8c-01c1f81dddff"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Dataset YOLO format"
      ],
      "metadata": {
        "id": "c648pGhlmw-z"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "import os\r\n",
        "from tqdm import tqdm\r\n",
        "import numpy as np\r\n",
        "import matplotlib.pyplot as plt\r\n",
        "from skimage import measure\r\n",
        "from matplotlib.patches import Rectangle\r\n",
        "from IPython.display import clear_output\r\n",
        "\r\n",
        "from utils.preprocess import *"
      ],
      "outputs": [],
      "metadata": {
        "id": "94Enuf8mPfPG"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "X_train = np.load('X_train_Seg_op.npy')\r\n",
        "Y_train = np.load('Y_train_Seg_op.npy')\r\n",
        "\r\n",
        "X_test = np.load('X_test_Seg_op.npy')\r\n",
        "Y_test = np.load('Y_test_Seg_op.npy') "
      ],
      "outputs": [],
      "metadata": {
        "id": "NNsJ7Z_a09xe"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "X_train = min_max_preprocessing(X_train)\r\n",
        "X_test  = min_max_preprocessing(X_test)\r\n",
        "\r\n",
        "X_train, mean, std = samplewise_preprocessing(X_train)\r\n",
        "X_test = featurewise_preprocessing(X_test, mean, std)"
      ],
      "outputs": [],
      "metadata": {
        "id": "nsvE4CJOp98T"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "for i in tqdm(range(X_train.shape[0])):\r\n",
        "  plt.imshow(X_train[i], cmap=\"gray\")\r\n",
        "  plt.axis('off')\r\n",
        "  plt.savefig(f'BoundingBox_dataset/Version_1/images/train/{i+1}.png', bbox_inches='tight', pad_inches=0)\r\n",
        "  plt.cla()\r\n",
        "\r\n",
        "for i in tqdm(range(X_test.shape[0])):\r\n",
        "  plt.imshow(X_test[i], cmap=\"gray\")\r\n",
        "  plt.axis('off')\r\n",
        "  plt.savefig(f'BoundingBox_dataset/Version_1/images/valid/{i+1}.png', bbox_inches='tight', pad_inches=0)\r\n",
        "  plt.cla()"
      ],
      "outputs": [],
      "metadata": {
        "id": "ZAEl1pBToln4"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "# Convert segmentation masks into YOLO labels using the Skimage library\r\n",
        "\r\n",
        "for i in tqdm(range(Y_train.shape[0])):\r\n",
        "  label_mask = measure.label(Y_train[i])\r\n",
        "  props = measure.regionprops(label_mask)\r\n",
        "  f = open(f'BoundingBox_dataset/Version_1/labels/train/{i+1}.txt', 'w')\r\n",
        "  l = []\r\n",
        "  for prop in props:\r\n",
        "    width = (prop.bbox[3]-prop.bbox[1])/Y_train.shape[1]\r\n",
        "    height = (prop.bbox[2]-prop.bbox[0])/Y_train.shape[1]\r\n",
        "    center_x = prop.centroid[1]/Y_train.shape[1]\r\n",
        "    center_y = prop.centroid[0]/Y_train.shape[1]\r\n",
        "    l.append(f'0 {center_x:.4f} {center_y:.4f} {width:.4f} {height:.4f}\\n')\r\n",
        "  f.writelines(l)\r\n",
        "  f.close\r\n",
        "\r\n",
        "for i in tqdm(range(Y_test.shape[0])):\r\n",
        "  label_mask = measure.label(Y_test[i])\r\n",
        "  props = measure.regionprops(label_mask)\r\n",
        "  f = open(f'BoundingBox_dataset/Version_1/labels/valid/{i+1}.txt', 'w')\r\n",
        "  l = []\r\n",
        "  for prop in props:\r\n",
        "    width = (prop.bbox[3]-prop.bbox[1])/Y_test.shape[1]\r\n",
        "    height = (prop.bbox[2]-prop.bbox[0])/Y_test.shape[1]\r\n",
        "    center_x = prop.centroid[1]/Y_test.shape[1]\r\n",
        "    center_y = prop.centroid[0]/Y_test.shape[1]\r\n",
        "    l.append(f'0 {center_x:.4f} {center_y:.4f} {width:.4f} {height:.4f}\\n')\r\n",
        "  f.writelines(l)\r\n",
        "  f.close"
      ],
      "outputs": [],
      "metadata": {
        "id": "-A3HuS5BonZP",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "f4f049cf-4ded-4541-a08e-03fa32dc511e"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "def box_union(l, idxs):\r\n",
        "\r\n",
        "  \"\"\"\r\n",
        "  Function to join boxes from a list in YOLO format.\r\n",
        "\r\n",
        "  Inputs\r\n",
        "  l: List with the bounding boxes in YOLO format\r\n",
        "  idxs: List of indexes of the boxes to join\r\n",
        "  \"\"\"\r\n",
        "  \r\n",
        "  l_ = []\r\n",
        "  for idx in idxs:\r\n",
        "    l_.append(l[idx])\r\n",
        "  mini = np.min(np.array(l_), axis=0)\r\n",
        "  maxi = np.max(np.array(l_), axis=0)\r\n",
        "  l.append([mini[0], mini[1], maxi[2], maxi[3]])\r\n",
        "  return l"
      ],
      "outputs": [],
      "metadata": {
        "id": "kC75GoHMj9Qk"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "#i = int(os.listdir('BoundingBox_dataset/Version_2/labels/train/')[-1][0])\r\n",
        "i = 8\r\n",
        "fig, ax = plt.subplots()\r\n",
        "while True:\r\n",
        "  print(i)\r\n",
        "  f = open(f'BoundingBox_dataset/Version_2/labels/train/{i+1}.txt', 'w')\r\n",
        "  label_mask = measure.label(Y_train[i])\r\n",
        "  fig, ax = plt.subplots()\r\n",
        "  ax.imshow(Y_train[i], cmap=plt.cm.gray)\r\n",
        "  props = measure.regionprops(label_mask)\r\n",
        "  l = []\r\n",
        "  for j, prop in enumerate(props):\r\n",
        "    minr, minc, maxr, maxc = prop.bbox\r\n",
        "    l.append([minr, minc, maxr, maxc])\r\n",
        "    bx = (minc, maxc, maxc, minc, minc)\r\n",
        "    by = (minr, minr, maxr, maxr, minr)\r\n",
        "    ax.plot(bx, by, '-', linewidth=2.5, label=str(j))\r\n",
        "  ax.legend()\r\n",
        "  plt.show()\r\n",
        "\r\n",
        "  v = input('Ingrese los números de las bounding boxes que desea unir: ')\r\n",
        "\r\n",
        "  if v=='salir':\r\n",
        "    break\r\n",
        "\r\n",
        "  if v!='':\r\n",
        "    if v == 'ok':\r\n",
        "      data= []\r\n",
        "      for prop in props:\r\n",
        "        width = (prop.bbox[3]-prop.bbox[1])/Y_train.shape[1]\r\n",
        "        height = (prop.bbox[2]-prop.bbox[0])/Y_train.shape[1]\r\n",
        "        center_x = prop.centroid[1]/Y_train.shape[1]\r\n",
        "        center_y = prop.centroid[0]/Y_train.shape[1]\r\n",
        "        data.append(f'0 {center_x:.4f} {center_y:.4f} {width:.4f} {height:.4f}\\n')\r\n",
        "      f.writelines(data)\r\n",
        "      f.close()\r\n",
        "      i+=1\r\n",
        "      continue\r\n",
        "    idxs = v.split('-')\r\n",
        "    print(idxs)\r\n",
        "    if len(idxs) > 1:\r\n",
        "      p = []\r\n",
        "      for idx in idxs:\r\n",
        "        ind = [int(a) for a in idx.split(' ')]\r\n",
        "        l = box_union(l, ind)\r\n",
        "        for ind1 in ind:\r\n",
        "          p.append(ind1)\r\n",
        "      for ind2 in sorted(p, reverse=True):\r\n",
        "        l.pop(ind2)\r\n",
        "    else:\r\n",
        "      ind = [int(a) for a in idxs[0].split(' ')]\r\n",
        "      l = box_union(l, ind)\r\n",
        "    ax.clear()\r\n",
        "    ax.imshow(Y_train[i], cmap=plt.cm.gray)\r\n",
        "    data= []\r\n",
        "    for j, prop in enumerate(l):\r\n",
        "      minr, minc, maxr, maxc = prop\r\n",
        "      bx = (minc, maxc, maxc, minc, minc)\r\n",
        "      by = (minr, minr, maxr, maxr, minr)\r\n",
        "      ax.plot(bx, by, '-', linewidth=2.5, label=str(j))\r\n",
        "      width = (maxc-minc)/Y_train.shape[1]\r\n",
        "      height = (maxr-minr)/Y_train.shape[1]\r\n",
        "      center_x = (maxc+minc)/Y_train.shape[1]\r\n",
        "      center_y = (maxr+minr)/Y_train.shape[1]\r\n",
        "      data.append(f'0 {center_x:.4f} {center_y:.4f} {width:.4f} {height:.4f}\\n')\r\n",
        "    f.writelines(data)\r\n",
        "    f.close()\r\n",
        "    plt.show()\r\n",
        "    input(\"Press Enter to continue...\")\r\n",
        "    i+=1\r\n",
        "    "
      ],
      "outputs": [],
      "metadata": {
        "id": "hV_KONCqBsWF"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# YOLO model"
      ],
      "metadata": {
        "id": "x_CEbn02m3oo"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Set-up"
      ],
      "metadata": {
        "id": "asIkTXFICCfV"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "# %cd '/content/drive/MyDrive'\r\n",
        "# !git clone https://github.com/ultralytics/yolov5  # clone repo\r\n",
        "%cd '/content/Drive/MyDrive/COVID_19/yolov5'\r\n",
        "!pip install -r requirements.txt  # install dependencies\r\n",
        "\r\n",
        "import torch \r\n",
        "print('Setup complete. Using torch %s %s' % (torch.__version__, torch.cuda.get_device_properties(0) if torch.cuda.is_available() else 'CPU'))"
      ],
      "outputs": [],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "mYCc0jaDloLI",
        "outputId": "be8306f6-e429-4110-c76d-a255af5ebbb4"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "with open('BoundingBox_dataset/Version_3/data.yaml', 'w') as file:\r\n",
        "  file.writelines(['train: BoundingBox_dataset/Version_3/images/train/ \\n',\r\n",
        "                   'val: BoundingBox_dataset/Version_3/images/valid/ \\n', \r\n",
        "                   'nc: 1 \\n',\r\n",
        "                   'names: [\"GGO\"] \\n'])"
      ],
      "outputs": [],
      "metadata": {
        "id": "W4v8mHnD8Bhd"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "%cat 'BoundingBox_dataset/Version_3/data.yaml'"
      ],
      "outputs": [],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zhoO8O1c7_MR",
        "outputId": "d954a46b-b9c6-4b81-902b-42371e283efd"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "import yaml\r\n",
        "with open('BoundingBox_dataset/data.yaml', 'r') as stream:\r\n",
        "    num_classes = str(yaml.safe_load(stream)['nc'])"
      ],
      "outputs": [],
      "metadata": {
        "id": "bZvjhTPUnNiZ"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "%cat 'yolov5/models/yolov5x.yaml'"
      ],
      "outputs": [],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oEX8zL1v3w2H",
        "outputId": "b7cd4f47-2f01-47c0-8356-52d04a60ccf3"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Custom config file"
      ],
      "metadata": {
        "id": "XdGluD13p4De"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "#customize iPython writefile so we can write variables\r\n",
        "from IPython.core.magic import register_line_cell_magic\r\n",
        "\r\n",
        "@register_line_cell_magic\r\n",
        "def writetemplate(line, cell):\r\n",
        "    with open(line, 'w') as f:\r\n",
        "        f.write(cell.format(**globals()))"
      ],
      "outputs": [],
      "metadata": {
        "id": "t14hhyqdmw6O"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "%%writetemplate yolov5/models/custom_yolov5s.yaml\r\n",
        "\r\n",
        "# parameters\r\n",
        "nc: {num_classes}  # number of classes\r\n",
        "depth_multiple: 0.33  # model depth multiple\r\n",
        "width_multiple: 0.50  # layer channel multiple\r\n",
        "\r\n",
        "# anchors\r\n",
        "anchors:\r\n",
        "  - [10,13, 16,30, 33,23]  # P3/8\r\n",
        "  - [30,61, 62,45, 59,119]  # P4/16\r\n",
        "  - [116,90, 156,198, 373,326]  # P5/32\r\n",
        "\r\n",
        "# YOLOv5 backbone\r\n",
        "backbone:\r\n",
        "  # [from, number, module, args]\r\n",
        "  [[-1, 1, Focus, [64, 3]],  # 0-P1/2\r\n",
        "   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4\r\n",
        "   [-1, 3, BottleneckCSP, [128]],\r\n",
        "   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8\r\n",
        "   [-1, 9, BottleneckCSP, [256]],\r\n",
        "   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16\r\n",
        "   [-1, 9, BottleneckCSP, [512]],\r\n",
        "   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32\r\n",
        "   [-1, 1, SPP, [1024, [5, 9, 13]]],\r\n",
        "   [-1, 3, BottleneckCSP, [1024, False]],  # 9\r\n",
        "  ]\r\n",
        "\r\n",
        "# YOLOv5 head\r\n",
        "head:\r\n",
        "  [[-1, 1, Conv, [512, 1, 1]],\r\n",
        "   [-1, 1, nn.Upsample, [None, 2, 'nearest']],\r\n",
        "   [[-1, 6], 1, Concat, [1]],  # cat backbone P4\r\n",
        "   [-1, 3, BottleneckCSP, [512, False]],  # 13\r\n",
        "\r\n",
        "   [-1, 1, Conv, [256, 1, 1]],\r\n",
        "   [-1, 1, nn.Upsample, [None, 2, 'nearest']],\r\n",
        "   [[-1, 4], 1, Concat, [1]],  # cat backbone P3\r\n",
        "   [-1, 3, BottleneckCSP, [256, False]],  # 17 (P3/8-small)\r\n",
        "\r\n",
        "   [-1, 1, Conv, [256, 3, 2]],\r\n",
        "   [[-1, 14], 1, Concat, [1]],  # cat head P4\r\n",
        "   [-1, 3, BottleneckCSP, [512, False]],  # 20 (P4/16-medium)\r\n",
        "\r\n",
        "   [-1, 1, Conv, [512, 3, 2]],\r\n",
        "   [[-1, 10], 1, Concat, [1]],  # cat head P5\r\n",
        "   [-1, 3, BottleneckCSP, [1024, False]],  # 23 (P5/32-large)\r\n",
        "\r\n",
        "   [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)\r\n",
        "  ]"
      ],
      "outputs": [],
      "metadata": {
        "id": "uDxebz13RdRA"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "%%writetemplate yolov5/models/custom_yolov5x.yaml\r\n",
        "\r\n",
        "# parameters\r\n",
        "nc: {num_classes}  # number of classes\r\n",
        "depth_multiple: 1.33  # model depth multiple\r\n",
        "width_multiple: 1.25  # layer channel multiple\r\n",
        "\r\n",
        "# anchors\r\n",
        "anchors:\r\n",
        "  - [10,13, 16,30, 33,23]  # P3/8\r\n",
        "  - [30,61, 62,45, 59,119]  # P4/16\r\n",
        "  - [116,90, 156,198, 373,326]  # P5/32\r\n",
        "\r\n",
        "# YOLOv5 backbone\r\n",
        "backbone:\r\n",
        "  # [from, number, module, args]\r\n",
        "  [[-1, 1, Focus, [64, 3]],  # 0-P1/2\r\n",
        "   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4\r\n",
        "   [-1, 3, C3, [128]],\r\n",
        "   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8\r\n",
        "   [-1, 9, C3, [256]],\r\n",
        "   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16\r\n",
        "   [-1, 9, C3, [512]],\r\n",
        "   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32\r\n",
        "   [-1, 1, SPP, [1024, [5, 9, 13]]],\r\n",
        "   [-1, 3, C3, [1024, False]],  # 9\r\n",
        "  ]\r\n",
        "\r\n",
        "# YOLOv5 head\r\n",
        "head:\r\n",
        "  [[-1, 1, Conv, [512, 1, 1]],\r\n",
        "   [-1, 1, nn.Upsample, [None, 2, 'nearest']],\r\n",
        "   [[-1, 6], 1, Concat, [1]],  # cat backbone P4\r\n",
        "   [-1, 3, C3, [512, False]],  # 13\r\n",
        "\r\n",
        "   [-1, 1, Conv, [256, 1, 1]],\r\n",
        "   [-1, 1, nn.Upsample, [None, 2, 'nearest']],\r\n",
        "   [[-1, 4], 1, Concat, [1]],  # cat backbone P3\r\n",
        "   [-1, 3, C3, [256, False]],  # 17 (P3/8-small)\r\n",
        "\r\n",
        "   [-1, 1, Conv, [256, 3, 2]],\r\n",
        "   [[-1, 14], 1, Concat, [1]],  # cat head P4\r\n",
        "   [-1, 3, C3, [512, False]],  # 20 (P4/16-medium)\r\n",
        "\r\n",
        "   [-1, 1, Conv, [512, 3, 2]],\r\n",
        "   [[-1, 10], 1, Concat, [1]],  # cat head P5\r\n",
        "   [-1, 3, C3, [1024, False]],  # 23 (P5/32-large)\r\n",
        "\r\n",
        "   [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)\r\n",
        "  ]"
      ],
      "outputs": [],
      "metadata": {
        "id": "2TWOz_f-39Z7"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Train"
      ],
      "metadata": {
        "id": "N0F_bxhQqAB4"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "!python train.py --img 224 --batch 128 --epochs 500 --data 'BoundingBox_dataset/Version_3/data.yaml' --weights yolov5x.pt\r\n",
        "#--cfg 'yolov5/models/custom_yolov5x.yaml' --weights 'yolov5/runs/train/exp/weights/best.pt'"
      ],
      "outputs": [],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "K-q3OvzLp86O",
        "outputId": "d048bfec-bd67-4bb7-ccc2-17d9e9cdd0f8"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "from IPython.display import Image\r\n",
        "print(\"GROUND TRUTH TRAINING DATA:\")\r\n",
        "Image(filename='yolov5/runs/train/exp4/test_batch0_labels.jpg', width=900)"
      ],
      "outputs": [],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 935
        },
        "id": "ADkk2zdbr5JE",
        "outputId": "fabc5ff7-8d5e-4034-b259-6963623dd614"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "Image(filename='yolov5/runs/train/exp4/test_batch0_pred.jpg', width=900)"
      ],
      "outputs": [],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 917
        },
        "id": "mlxVISUUr5js",
        "outputId": "8624c4d7-6e26-414c-e2b7-cfe642696eac"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Test"
      ],
      "metadata": {
        "id": "dv17ZUa5bv73"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "!python test.py --img 224 --batch 128 --data 'BoundingBox_dataset/Version_3/data.yaml' --weights 'yolov5/runs/train/exp5/weights/best.pt'"
      ],
      "outputs": [],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bxmYn0d3bxGE",
        "outputId": "724f8487-eea0-4964-8a7f-57cec1f033fd"
      }
    }
  ]
}