tutorials/detection/cars_yolov7/car_detection__prepare_and_train.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Car detecion from aerial ortophoto with YOLOv7\n",
    "\n",
    "This notebook is just a code suplementary for `Deepness` QGIS plugin documentation.\n",
    "Please visit the documentation to see how easily use the model in QGIS: https://qgis-plugin-deepness.readthedocs.io/\n",
    "\n",
    "\n",
    "Model trained with dataset ITCVD ( https://arxiv.org/pdf/1801.07339.pdf )\n",
    "Spatial resolution of model: 10 cm/pixel.\n",
    "\n",
    "In this notebook we go though the training process, including data preprocessing and model export."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download yolov7\n",
    "Download yolov7 repository, install requirements, and put this notebook in the root directory of the repository:\n",
    "\n",
    "``` \n",
    "git clone https://github.com/WongKinYiu/yolov7 \n",
    "pip install -r yolov7/requirements.txt\n",
    "pip install scipy onnx\n",
    "cp $THIS_NOTEBOOK ./yolov7 \n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import scipy.io\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset manual download\n",
    "\n",
    "Download manually ITCVD dataset.\n",
    "\n",
    "At the moment of writing available at https://easy.dans.knaw.nl/ui/datasets/id/easy-dataset:157920/tab/2\n",
    "\n",
    "Create subdirectory for our data:\n",
    "```\n",
    "mkdir -p car/itcvd\n",
    "```\n",
    "\n",
    "Put the content of downloaded data `ITCVD/ITC_VD_Training_Testing_set/Training` in `./car/itcvd` (So that we will have a directory `./car/itcvd/GT` and `./car/itcvd/Images`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BASE_INPUT_DATA_DIR = 'car/itcvd/'\n",
    "BASE_OUTPUT_DATA_DIR = 'car/'\n",
    "\n",
    "img_output_dir = os.path.join(BASE_OUTPUT_DATA_DIR, 'images', 'all')\n",
    "label_output_dir = os.path.join(BASE_OUTPUT_DATA_DIR, 'labels', 'all')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data preprocessing and preparation\n",
    "\n",
    "The dataset images and labels are not suitable for yolov7. The images need to be cut into smaller parts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subimg_number = 0\n",
    "\n",
    "img_overlap = 0.05\n",
    "target_img_size = 640  # both dimensions equal\n",
    "stride = int(target_img_size * (1 - img_overlap))\n",
    "\n",
    "\n",
    "os.makedirs(img_output_dir, exist_ok=True)\n",
    "os.makedirs(label_output_dir, exist_ok=True)\n",
    "\n",
    "input_img_dir = os.path.join(BASE_INPUT_DATA_DIR, 'Image/')\n",
    "all_output_img_names = []\n",
    "\n",
    "for file_name in sorted(os.listdir(input_img_dir)):\n",
    "    if not file_name.endswith('.jpg'):\n",
    "        continue\n",
    "\n",
    "    input_image_number = file_name.strip('.jpeg')\n",
    "    mat_file_name = input_image_number + '.mat'\n",
    "    labels_mat_path = os.path.join(BASE_INPUT_DATA_DIR, f'GT/{mat_file_name}')\n",
    "    img_path = os.path.join(input_img_dir, f'{file_name}')\n",
    "\n",
    "    labels_mat = scipy.io.loadmat(labels_mat_path)\n",
    "    full_img = Image.open(img_path)\n",
    "    full_img_numpy = np.array(full_img)\n",
    "\n",
    "    car_labels = labels_mat[f'x{input_image_number}']\n",
    "\n",
    "    input_img_size = full_img_numpy.shape\n",
    "    bins_x = (input_img_size[1] - target_img_size) // stride + 1\n",
    "    bins_y = (input_img_size[0] - target_img_size) // stride + 1\n",
    "\n",
    "    for i in range(bins_x):\n",
    "        for j in range(bins_y):\n",
    "            x_start = i * stride\n",
    "            y_start = j * stride\n",
    "            x_end = x_start + target_img_size\n",
    "            y_end = y_start + target_img_size\n",
    "\n",
    "            if subimg_number == 1:\n",
    "                a = 1\n",
    "\n",
    "            yolo_labels = []\n",
    "            for car_label in car_labels:\n",
    "                x_left, y_upper, x_right, y_bottom, _, _ = car_label\n",
    "                x_centre, y_centre = (x_left + x_right) / 2, (y_upper + y_bottom) / 2\n",
    "                if x_centre > x_start and x_centre < x_end and y_centre < y_end and y_centre > y_start:\n",
    "                    x_relative_centre = (x_centre - x_start) / target_img_size\n",
    "                    y_relative_centre = (y_centre - y_start) / target_img_size\n",
    "                    x_relative_width = (x_right - x_left) / target_img_size\n",
    "                    y_relative_width = (y_bottom - y_upper) / target_img_size\n",
    "                    class_id = 0\n",
    "                    # yolo_label = [class_id, x_relative_centre, y_relative_centre, x_relative_width, y_relative_width]\n",
    "                    yolo_label = [class_id, x_relative_centre, y_relative_centre, x_relative_width, y_relative_width]\n",
    "                    yolo_labels.append(list(map(str, yolo_label)))\n",
    "\n",
    "            img_slice = full_img_numpy[y_start:y_end, x_start:x_end]\n",
    "            \n",
    "            file_base_name = 's' + str(subimg_number).zfill(5)\n",
    "            img_file_path = os.path.join(img_output_dir, f'{file_base_name}.jpg')\n",
    "            subimg_number += 1\n",
    "            im = Image.fromarray(img_slice)\n",
    "            im.save(img_file_path)\n",
    "            all_output_img_names.append(f'{file_base_name}.jpg')\n",
    "\n",
    "            label_file_path = os.path.join(label_output_dir, f'{file_base_name}.txt')\n",
    "            txt = '\\n'.join([' '.join(label) for label in yolo_labels])\n",
    "            with open(label_file_path, 'wt') as file:\n",
    "                file.write(txt)            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create configuration files for yolo, with list of training/testing files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images_numbers = list(range(subimg_number))\n",
    "\n",
    "x = np.array(all_output_img_names)\n",
    "N = len(x)\n",
    "np.random.seed(44)\n",
    "indices = np.random.permutation(N)\n",
    "training_idx, test_idx, val_idx = indices[:int(0.8*N)], indices[int(0.8*N):int(0.9*N)], indices[int(0.9*N):]\n",
    "training, test, val = x[training_idx], x[test_idx], x[val_idx]\n",
    "\n",
    "\n",
    "def create_as_images_listing(set_name, file_numbers):\n",
    "    txts = []\n",
    "    for file_number in file_numbers:\n",
    "        img_file_name = f'{file_number}'\n",
    "        txts.append(f'./images/all/{img_file_name}') \n",
    "\n",
    "    txt = '\\n'.join(txts)\n",
    "    txt_file_path = os.path.join(BASE_OUTPUT_DATA_DIR, set_name + '.txt')\n",
    "    with open(txt_file_path, 'wt') as file:\n",
    "        file.write(txt) \n",
    "        \n",
    "\n",
    "create_as_images_listing('train', sorted(training))\n",
    "create_as_images_listing('test', sorted(test))\n",
    "create_as_images_listing('val', sorted(val))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CAR_YAML_CONENT = \"\"\"\n",
    "# CAR\n",
    "\n",
    "# download command/URL (optional)\n",
    "download: echo \"no download\"\n",
    "\n",
    "# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]\n",
    "train: ./car/train.txt\n",
    "val: ./car/val.txt\n",
    "test: ./car/test.txt\n",
    "\n",
    "# number of classes\n",
    "nc: 1\n",
    "\n",
    "# class names\n",
    "names: [\"car\"]\n",
    "\"\"\"\n",
    "\n",
    "with open('./car/car.yaml', 'wt') as file:\n",
    "    file.write(CAR_YAML_CONENT) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Run the training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python train.py \\\n",
    "    --workers 8 \\\n",
    "    --device 0 \\\n",
    "    --batch-size 2 \\\n",
    "    --data car/car.yaml \\\n",
    "    --img 640 640 \\\n",
    "    --cfg cfg/training/yolov7-tiny.yaml \\\n",
    "    --weights yolov7-tiny.pt \\\n",
    "    --name yolov7-car-detector \\\n",
    "    --hyp data/hyp.scratch.custom.yaml \\\n",
    "    --epochs 15 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python test.py  \\\n",
    "    --device 0 \\\n",
    "    --batch-size 2 \\ \n",
    "    --data car/car.yaml \\\n",
    "    --img 640  \\\n",
    "    --weights runs/train/yolov7-car-detector/weights/best.pt_xx \\\n",
    "    --name yolov7-car-detector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Export the model to ONNX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python export.py \\\n",
    "    --weights runs/train/yolov7-car-detector/weights/best.pt \\\n",
    "    --grid \\\n",
    "    --simplify \\\n",
    "    --img-size 640 640"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Add metadata for `Deepness` plugin to run the model smoothly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import onnx\n",
    "\n",
    "model = onnx.load('runs/train/yolov7-car-detector3/weights/best.onnx')\n",
    "\n",
    "class_names = {\n",
    "    0: 'car',\n",
    "}\n",
    "\n",
    "m1 = model.metadata_props.add()\n",
    "m1.key = 'model_type'\n",
    "m1.value = json.dumps('Detector')\n",
    "\n",
    "m2 = model.metadata_props.add()\n",
    "m2.key = 'class_names'\n",
    "m2.value = json.dumps(class_names)\n",
    "\n",
    "m3 = model.metadata_props.add()\n",
    "m3.key = 'resolution'\n",
    "m3.value = json.dumps(10)\n",
    "\n",
    "m4 = model.metadata_props.add()\n",
    "m4.key = 'tiles_overlap'\n",
    "m4.value = json.dumps(10)\n",
    "\n",
    "m4 = model.metadata_props.add()\n",
    "m4.key = 'det_conf'\n",
    "m4.value = json.dumps(0.3)\n",
    "\n",
    "m4 = model.metadata_props.add()\n",
    "m4.key = 'det_iou_thresh'\n",
    "m4.value = json.dumps(0.7)\n",
    "\n",
    "\n",
    "FIANL_MODEL_FILE_PATH = os.path.abspath('runs/train/yolov7-car-detector2/weights/car_aerial_detection_yolo7_ITCVD_deepness.onnx')\n",
    "onnx.save(model, FIANL_MODEL_FILE_PATH)\n",
    "\n",
    "print(f'Your ONNX model with metadata is at: {FIANL_MODEL_FILE_PATH}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Now you can load the model in Deepness!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.6"
  },
  "vscode": {
   "interpreter": {
    "hash": "57069042d30620f0a3856c40b208c9eb0a4618b3f277ee0b3ee7d5bf24466b9e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}