BLACKPINK_detection_with_YOLOv5.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "BLACKPINK detection.ipynb",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e-BTJICyPzEG"
      },
      "source": [
        "# Table of Contents\r\n",
        "1. train, val.txt 파일 만들기\r\n",
        "2. data.yaml 파일 수정하기\r\n",
        "3. 사용할 모델의 yaml파일 수정하기\r\n",
        "4. YOLO v5 train\r\n",
        "5. YOLO v5 detection"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZgM5rZEBduuG"
      },
      "source": [
        "# 구글코랩 사용 시 설치 필요 (세션 시작때마다 해주어야 함)\r\n",
        "!pip install -U PyYAML"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tmD1ugmbevIX"
      },
      "source": [
        "# 필요모듈 import\r\n",
        "import numpy as np\r\n",
        "import pandas as pd\r\n",
        "from glob import glob\r\n",
        "import json\r\n",
        "import os\r\n",
        "import yaml\r\n",
        "import random\r\n",
        "import numpy as np\r\n",
        "import shutil\r\n",
        "import cv2\r\n",
        "import imgaug as ia\r\n",
        "import imgaug.augmenters as iaa\r\n",
        "from imgaug.augmentables.bbs import BoundingBox, BoundingBoxesOnImage\r\n",
        "from cv2 import imshow\r\n",
        "from IPython.display import display\r\n",
        "from google.colab.patches import cv2_imshow\r\n",
        "import cv2\r\n",
        "import json"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Rd9-emQ4lTOJ"
      },
      "source": [
        "## 1. train, val text 파일 만들기\r\n",
        "- train, validation 이미지를 모두 train에 사용\r\n",
        "    - 모델 구동을 위해 100번째 이미지부터 validation으로 설정하여 훈련시킴\r\n",
        "- test에는 앵콜무대 영상을 사용함"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P9rEIT5s3Zr1"
      },
      "source": [
        "train_ls = glob('/content/drive/MyDrive/Data/blackpink/origin_data/*.jpg')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZXukWvGMBgL1"
      },
      "source": [
        "val_ls = glob('/content/drive/MyDrive/Data/blackpink/origin_data/*.jpg')[100:]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ecKdqrHKD5JF"
      },
      "source": [
        "with open('/content/drive/MyDrive/Data/blackpink/train.txt', 'w') as f:\r\n",
        "    f.write('\\n'.join(train_ls) + '\\n')\r\n",
        "\r\n",
        "with open('/content/drive/MyDrive/Data/blackpink/val.txt', 'w') as f:\r\n",
        "    f.write('\\n'.join(val_ls) + '\\n')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Prhu8zdqlaRz"
      },
      "source": [
        "## 2. data.yaml 파일 수정하기"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "e6IZl9o7lcjH"
      },
      "source": [
        "# yaml파일 열기\r\n",
        "with open ('/content/drive/MyDrive/Data/blackpink/data.yaml','r') as f:\r\n",
        "    data = yaml.load(f)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "77hd7Hle38nT"
      },
      "source": [
        "print(data)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kRdYx7eLlh_s"
      },
      "source": [
        "data['names'] = ['Rose', 'Lisa', 'Jisu', 'Jenny' ]\r\n",
        "data['nc'] = 4\r\n",
        "data['train'] = '/content/drive/MyDrive/Data/blackpink/train.txt'\r\n",
        "data['val'] = '/content/drive/MyDrive/Data/blackpink/val.txt'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z5yu9h1NCUeJ"
      },
      "source": [
        "data"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8gshPgtbl6sK"
      },
      "source": [
        "with open('/content/drive/MyDrive/Data/blackpink/data.yaml', 'w') as f:\r\n",
        "    yaml.dump(data, f)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k07s5BgL6Q8c"
      },
      "source": [
        "## 3. 사용할 모델의 yaml파일 수정하기\r\n",
        "- './yolov5/models/yolov5s.yaml' yaml파일 열기 (구글코랩에서는 바로 열고 수정 가능/ 로컬에서 사용 시, 위의 방법처럼 with open으로 열어서 수정필요)\r\n",
        "- parameters에 nc 4로 수정 (실제 detection 하고자 하는 class의 수로 변경)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "245x2bBEmanp"
      },
      "source": [
        "## 4. YOLO v5 train"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-TPGj_KbmZWn"
      },
      "source": [
        "cd /content/drive/MyDrive/yolov5"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VWYJkeqnQFxA"
      },
      "source": [
        "# 전체데이터를 train 시키기\r\n",
        "!python train.py --data '/content/drive/MyDrive/Data/blackpink/data.yaml' --epochs 1000 --cfg '/content/drive/MyDrive/yolov5/models/yolov5s.yaml' --weights '' --batch-size 64 --name blackpink_detection"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X0AtMjznOgbJ"
      },
      "source": [
        "## 5. YOLO v5 detection"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nkOX0KV9M7f6"
      },
      "source": [
        "# encore 영상에 detection test\r\n",
        "!python detect.py --weights \"/content/drive/MyDrive/yolov5/runs/train/blackpink_detection/weights/best.pt\" --conf 0.3 --source \"/content/drive/MyDrive/Dataset/blackpink_encore.mp4\""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}