H_Tatarstan2022_Final.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "H_Tatarstan2022_Final.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "synyZ9WcZU8_"
      ],
      "authorship_tag": "ABX9TyOrZk/h3n7PJQa+QlDzgWWZ",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/TayJen/H_Tatarsan2022/blob/main/H_Tatarstan2022_Final.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Импорт библиотек + необходимые установки"
      ],
      "metadata": {
        "id": "synyZ9WcZU8_"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cQUF6r2oH4iC"
      },
      "outputs": [],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Здесь необходимо вставить путь, по которому будут располагаться все файлы, включая данные, модель и аннотации.\n",
        "\n"
      ],
      "metadata": {
        "id": "26MK6sLOZbD1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content/drive/Shareddrives/data_drive/H_Tatarstan2022\n",
        "!pwd"
      ],
      "metadata": {
        "id": "NdkbsNv8IDR5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Также в данной папке необходимо создать папку `data`. Внутри нее разместить папку `zipped`, в которой будут находиться используемые данные (`train_dataset_train.zip` и `test_dataset_test.zip`)\n",
        "\n",
        "Таким образом полный путь изначальных данных будет:\n",
        "\n",
        "*   `./data/zipped/train_dataset_train.zip`\n",
        "*   `./data/zipped/test_dataset_test.zip`"
      ],
      "metadata": {
        "id": "rYc3AmwGaiQ6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!unzip data/zipped/train_dataset_train.zip -d data/\n",
        "!unzip data/zipped/test_dataset_test.zip -d data/"
      ],
      "metadata": {
        "id": "AdLzrDWPIDT9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Устанавливаем необходимое для модели (модель - Yolov7)"
      ],
      "metadata": {
        "id": "y4Wb0ejeamv8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!git clone https://github.com/WongKinYiu/yolov7.git\n",
        "!pip install --quiet -r yolov7/requirements.txt\n",
        "!pip install --quiet map-boxes"
      ],
      "metadata": {
        "id": "3pxxbvfOIDWT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import pandas as pd\n",
        "import numpy as np\n",
        "\n",
        "np.random.seed(59)"
      ],
      "metadata": {
        "id": "jHvAWBLiarb4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Подготовка данных"
      ],
      "metadata": {
        "id": "Mb73RUdyauU5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Создадим обучающий датасет в формате yolo\n",
        "!mkdir ./data_for_yolo\n",
        "!mkdir ./data_for_yolo/data\n",
        "!mkdir ./data_for_yolo/data/images\n",
        "!mkdir ./data_for_yolo/data/labels\n",
        "!mkdir ./data_for_yolo/data/images/train\n",
        "!mkdir ./data_for_yolo/data/labels/train\n",
        "!mkdir ./data_for_yolo/data/images/test\n",
        "!mkdir ./data_for_yolo/data/labels/test\n",
        "\n",
        "# Для создания аннотаций в нужном формате\n",
        "!mkdir ./data/train/yolo_labels"
      ],
      "metadata": {
        "id": "u-Ke3v5Bareo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import yaml\n",
        "\n",
        "# Последняя версия модели Yolo - Yolov7 немного привередлива и требует абсолютные пути\n",
        "# до данных вместо относительных\n",
        "def create_yaml_cont():\n",
        "    \"\"\"\n",
        "        train: /content/drive/Shareddrives/data_drive/H_Tatarstan2022/data_for_yolo/data/images/train/\n",
        "        val: /content/drive/Shareddrives/data_drive/H_Tatarstan2022/data_for_yolo/data/images/test/\n",
        "\n",
        "        # number of classes\n",
        "        nc: 5\n",
        "\n",
        "        # class names\n",
        "        names: ['human', 'head', 'face', 'car', 'carplate']\n",
        "    \"\"\"\n",
        "\n",
        "    yaml_content = {\n",
        "        'train': '/content/drive/Shareddrives/data_drive/H_Tatarstan2022/data_for_yolo/data/images/train/',\n",
        "        'val': '/content/drive/Shareddrives/data_drive/H_Tatarstan2022/data_for_yolo/data/images/test',\n",
        "        'nc': 5,\n",
        "        'names': ['human', 'head', 'face', 'car', 'carplate']\n",
        "    }\n",
        "\n",
        "    with open('data_for_yolo/dataset.yaml', 'w') as f_yml:\n",
        "        yaml.dump(yaml_content, f_yml)\n",
        "\n",
        "\n",
        "create_yaml_cont()"
      ],
      "metadata": {
        "id": "3QR_pDu6IDYR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Сделаем разметку как в Yolo, после этого нужным образом перемешаем и\n",
        "# закинем в папку data_for_yolo\n",
        "from tqdm import tqdm\n",
        "\n",
        "\n",
        "name_to_class_num = {\n",
        "    'human': 0,\n",
        "    'head': 1,\n",
        "    'face': 2,\n",
        "    'car': 3,\n",
        "    'carplate': 4\n",
        "}\n",
        "\n",
        "def yolo_annotation():\n",
        "    old_path = './data/train/labels/'\n",
        "    new_path = './data/train/yolo_labels/'\n",
        "\n",
        "    for img_name in tqdm(os.listdir('data/train/images/')):\n",
        "        img_name = img_name.split('.')[0]\n",
        "        with open(new_path + img_name + '.txt', 'w') as f_new:\n",
        "            for f in os.listdir(old_path):\n",
        "                if f.startswith(img_name):\n",
        "                    class_name = f.split('.')[0].split('_')[-1]\n",
        "                    class_num = name_to_class_num[class_name]\n",
        "                    with open(old_path + f, 'r') as f_old:\n",
        "                        for row in f_old.readlines():\n",
        "                            f_new.write(str(class_num) + row[1:])\n",
        "\n",
        "\n",
        "yolo_annotation()"
      ],
      "metadata": {
        "id": "QG97olDGIDaU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from sklearn.model_selection import train_test_split\n",
        "\n",
        "\n",
        "# Используем 20% датасета для валидации\n",
        "images_names = os.listdir('./data/train/images')\n",
        "train_images, test_images = train_test_split(images_names, test_size=0.2,\n",
        "                                             shuffle=True, random_state=59)\n",
        "\n",
        "len(train_images), len(test_images)"
      ],
      "metadata": {
        "id": "GHQ4pjeSIDcg"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_images = set(train_images)\n",
        "test_images = set(test_images)\n",
        "\n",
        "len(train_images), len(test_images)"
      ],
      "metadata": {
        "id": "cZ_VgNQSIDel"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from shutil import copy\n",
        "from tqdm import tqdm\n",
        "\n",
        "\n",
        "def copy_to_yolo():\n",
        "    old_path_img = './data/train/images/'\n",
        "    old_path_lbl = './data/train/yolo_labels/'\n",
        "    yolo_path_img = './data_for_yolo/data/images/'\n",
        "    yolo_path_lbl = './data_for_yolo/data/labels/'\n",
        "\n",
        "    for image_name in tqdm(os.listdir(old_path_img)):\n",
        "        lbl_name = image_name.split('.')[0] + '.txt'\n",
        "        # print(image_name, lbl_name)\n",
        "        if image_name in train_images:\n",
        "            copy(old_path_img + image_name, yolo_path_img + 'train/' + image_name)\n",
        "            copy(old_path_lbl + lbl_name, yolo_path_lbl + 'train/' + lbl_name)\n",
        "        elif image_name in test_images:\n",
        "            copy(old_path_img + image_name, yolo_path_img + 'test/' + image_name)\n",
        "            copy(old_path_lbl + lbl_name, yolo_path_lbl + 'test/' + lbl_name)\n",
        "\n",
        "\n",
        "copy_to_yolo()"
      ],
      "metadata": {
        "id": "vbKEZ4tfIDgk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "len(os.listdir('./data_for_yolo/data/images/train')), len(os.listdir('./data_for_yolo/data/images/test'))"
      ],
      "metadata": {
        "id": "Ud9ecPUhIDin"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "len(os.listdir('./data_for_yolo/data/labels/train')), len(os.listdir('./data_for_yolo/data/labels/test'))"
      ],
      "metadata": {
        "id": "6SL8QS4fIDk9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Модель"
      ],
      "metadata": {
        "id": "vjIRGAFwbjl6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!nvidia-smi"
      ],
      "metadata": {
        "id": "lzDTAo1XIDnP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "path_to_data = './data_for_yolo/dataset.yaml'"
      ],
      "metadata": {
        "id": "iAUavlyVIDpT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%load_ext tensorboard\n",
        "%tensorboard --logdir ./runs/train"
      ],
      "metadata": {
        "id": "6gkF3jy1IDr0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "В результате многочисленных экспериментов лучший результат показала модель с разрешением 1280 (`batch_size` равен 2, так как бесплатная GPU от Google Colab'a не вмещает в себя больше).\n",
        "\n",
        "Также возможно лучшее разрешение показало бы лучшие результаты, но в бесплатной версии обучение идет слишком долго и память слишком мала."
      ],
      "metadata": {
        "id": "7BmRFJUCb72C"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "!!! Внимание, чтобы обучение модели запустилось без ошибок, на этом этапе необходимо в созданную папку yolov7 загрузить веса, скачанные с официального репозитория модели: [YOLOv7](https://github.com/WongKinYiu/yolov7)\n",
        "\n",
        "В разделе `Transfer Learning` имеются веса предобученных моделей различной сложности. В данном решении используются веса средней по сложности и памяти модели `yolov7x_training.pt`\n",
        "\n",
        "После скачивания весов их необходимо поместить в папку модели, если этого не сделать то будет выдаваться ошибка об отсутствии возможности скачать веса"
      ],
      "metadata": {
        "id": "x6QqZVuWcZdn"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!python ./yolov7/train.py --img 1280 --batch 2 --epochs 35 --freeze 10 --data {path_to_data} --weights ./yolov7/yolov7x_training.pt --name yolov7x_general_1280_results --cache"
      ],
      "metadata": {
        "id": "amJj1NzlIDuA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Модель обучается около 4-5 часов, если не следить за ней то она может вылететь"
      ],
      "metadata": {
        "id": "pox_FP8HdOsp"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Предикт для тестовой части\n",
        "\n",
        "По окончанию обучения модели веса лучшей модели за все обучение будут сохранены в папке которая была указана при запуске обучения, в разделе `weights` под названием `best.pt`.\n",
        "\n",
        "Лучший порог `--conf` оказался равным 0.01"
      ],
      "metadata": {
        "id": "oDJxYBvuda7K"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!python3 ./yolov7/detect.py --weights ./runs/train/yolov7x_general_1280_results/weights/best.pt --img 1280 --conf 0.01 --source ./data_for_yolo/data/images/test --save-txt --save-conf"
      ],
      "metadata": {
        "id": "e_GBucxXbxHL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Функции для получения меток и границ объектов (для теста еще и уверенности) из текстовых файлов в списки"
      ],
      "metadata": {
        "id": "O-PpHhUQd64N"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def get_soluton_labels(path_to_txt_folder):\n",
        "    solutions = []\n",
        "    for detection_file in os.listdir(path_to_txt_folder):\n",
        "        img_name = detection_file.split('.')[0] + '.jpg'\n",
        "        with open(path_to_txt_folder + detection_file, 'r') as f:\n",
        "            data = f.read()\n",
        "            data = [i for i in data.split('\\n') if i != '']\n",
        "        for line in data:\n",
        "            val = [float(i) for i in line.split()]\n",
        "            cls, xywh, conf = val[0], val[1:5], val[5]\n",
        "            center_x, center_y, width, height = xywh\n",
        "            xmin = center_x - (width / 2)\n",
        "            xmax = center_x + (width / 2)\n",
        "            ymin = center_y - (height / 2)\n",
        "            ymax = center_y + (height / 2)\n",
        "            solutions.append([img_name, int(cls), conf, xmin, xmax, ymin, ymax])\n",
        "\n",
        "    return solutions\n",
        "\n",
        "def get_test_labels(path_to_txt_folder):\n",
        "    solutions = []\n",
        "    for detection_file in os.listdir(path_to_txt_folder):\n",
        "        img_name = detection_file.split('.')[0] + '.jpg'\n",
        "        with open(path_to_txt_folder + detection_file, 'r') as f:\n",
        "            data = f.read()\n",
        "            data = [i for i in data.split('\\n') if i != '']\n",
        "        for line in data:\n",
        "            val = [float(i) for i in line.split()]\n",
        "            cls, center_x, center_y, width, height = val\n",
        "            xmin = center_x - (width / 2)\n",
        "            xmax = center_x + (width / 2)\n",
        "            ymin = center_y - (height / 2)\n",
        "            ymax = center_y + (height / 2)\n",
        "            solutions.append([img_name, int(cls), xmin, xmax, ymin, ymax])\n",
        "    return solutions"
      ],
      "metadata": {
        "id": "dav0LJi4bxJM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "solution_v7x = get_soluton_labels('./runs/detect/exp/labels/')\n",
        "solution_v7x_df = pd.DataFrame(solution_v7x, columns=['ImageID', 'LabelName', 'Conf', 'XMin', 'XMax', 'YMin', 'YMax'])\n",
        "solution_v7x_df.head()"
      ],
      "metadata": {
        "id": "6lOg1CJDbxLT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "test_labels = get_test_labels('./data_for_yolo/data/labels/test/')\n",
        "test_labels_df = pd.DataFrame(test_labels, columns=['ImageID', 'LabelName', 'XMin', 'XMax', 'YMin', 'YMax'])\n",
        "test_labels_df.head()"
      ],
      "metadata": {
        "id": "zRV2yNs5bxNd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Проверим скор на отложенной выборке"
      ],
      "metadata": {
        "id": "hnr79avFeM2d"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from map_boxes import mean_average_precision_for_boxes as map_boxes_fn\n",
        "\n",
        "mean_ap, average_precisions = map_boxes_fn(test_labels, solution_v7x_df,\n",
        "                                           iou_threshold=0.25, verbose=False)\n",
        "mean_ap"
      ],
      "metadata": {
        "id": "8S7xa2xBbxPS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Предсказание для теста"
      ],
      "metadata": {
        "id": "9QwjWry7eQgG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!python3 ./yolov7/detect.py --weights ./runs/train/yolov7x_general_1280_results/weights/best.pt --img 1280 --conf 0.01 --source ./data/test/images --save-txt --save-conf"
      ],
      "metadata": {
        "id": "7KJwNTeqePC8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "solution_site_yolov7 = get_soluton_labels('./runs/detect/exp2/labels/')\n",
        "solution_site_yolov7_df = pd.DataFrame(solution_site_yolov7, columns=['ImageID', 'LabelName', 'Conf', 'XMin', 'XMax', 'YMin', 'YMax'])\n",
        "solution_site_yolov7_df.head()"
      ],
      "metadata": {
        "id": "ViLMAEbrePE-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Проверяем что нет никаких глупых ошибок"
      ],
      "metadata": {
        "id": "j-oywpGkecv1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "solution_site_yolov7_df.describe()"
      ],
      "metadata": {
        "id": "Wu3aRARQePHB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "При отправке решений необходимо соблюдать следующее соответствие классов:\n",
        "\n",
        "`'car': 0, 'head': 1, 'face': 2, 'human': 3, 'carplate': 4`"
      ],
      "metadata": {
        "id": "CMS4CCdRfJFY"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Поменяем свой формат меток на необходимый"
      ],
      "metadata": {
        "id": "0R1K7PmpfT6C"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "my_to_site = {\n",
        "    0: 3,\n",
        "    1: 1,\n",
        "    2: 2,\n",
        "    3: 0,\n",
        "    4: 4\n",
        "}\n",
        "\n",
        "solution_site_yolov7_df['LabelName'] = solution_site_yolov7_df.LabelName.apply(lambda x: my_to_site[x])"
      ],
      "metadata": {
        "id": "2tp4zjEaePI3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!mkdir solutions"
      ],
      "metadata": {
        "id": "HfydtJI-ePKr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "solution_site_yolov7_df.to_csv(\"./solutions/yolov7x_general_1280_results_conf01.csv\",\n",
        "                               sep=';', index=False)"
      ],
      "metadata": {
        "id": "L5s3lMXufbX8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Проверка предсказания\n",
        "\n",
        "Удостоверимся что файл сохранился"
      ],
      "metadata": {
        "id": "XbuzLrHafvOc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "path = \"./solutions/yolov7x_general_1280_results_conf01.csv\"\n",
        "\n",
        "df = pd.read_csv(path, sep=';')\n",
        "df.head()"
      ],
      "metadata": {
        "id": "2Ldtkej4fbZ-"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}