YOLOX_Colaboratory_Training_Sample.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "YOLOX_Colaboratory_Training_Sample.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KndWvALzfoG_"
      },
      "source": [
        "# YOLOX 依存パッケージインストール(YOLOX Dependent Package Install)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YpviPxKHfh59"
      },
      "source": [
        "!git clone https://github.com/Megvii-BaseDetection/YOLOX -b 0.3.0"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "maY7_U3gLhQA"
      },
      "source": [
        "%cd YOLOX\n",
        "\n",
        "!pip install -U pip && pip install -r requirements.txt\n",
        "!pip install -v -e .  "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install loguru thop ninja onnx onnxsim onnxruntime"
      ],
      "metadata": {
        "id": "KslVK6hveoHw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k1t1Hb74iIZo"
      },
      "source": [
        "# PyCocoToolsインストール(PyCocoTools Install)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YCje1A8mhtLy"
      },
      "source": [
        "!pip install cython\n",
        "!pip install 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iWp26peevulP"
      },
      "source": [
        "# データセットダウンロード(Download Dataset)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "frzsMeeetO1y"
      },
      "source": [
        "%cd /content/\n",
        "\n",
        "use_sample_image = True\n",
        "\n",
        "if use_sample_image:\n",
        "    !git clone https://github.com/Kazuhito00/YOLOX-Colaboratory-Training-Sample.git"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NC3Frlnzz5eC"
      },
      "source": [
        "# 学習/検証データ分割(Train/Validation split data)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pkp7yRJPv0_1"
      },
      "source": [
        "import os\n",
        "\n",
        "# 独自のデータを使用する場合は、パスを指定してください\n",
        "# Please fill in the path if you want to use your own data\n",
        "if use_sample_image:\n",
        "    dataset_directory = 'YOLOX-Colaboratory-Training-Sample/02.annotation_data'\n",
        "else:\n",
        "    dataset_directory = ''\n",
        "\n",
        "# 学習/検証データパス(train/validation data path)\n",
        "train_directory = './train'\n",
        "validation_directory = './validation'\n",
        "\n",
        "# 学習データ格納ディレクトリ作成(Create training data storage directory)\n",
        "os.makedirs(train_directory, exist_ok=True)\n",
        "# 検証データ格納ディレクトリ作成(Create verification data storage directory)\n",
        "os.makedirs(validation_directory, exist_ok=True)"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BJieAq9IywqQ"
      },
      "source": [
        "import glob\n",
        "import shutil\n",
        "import random\n",
        "\n",
        "# 学習データの割合(Percentage of training data)\n",
        "train_ratio = 0.8\n",
        "\n",
        "# コピー元ファイルリスト取得(Get copy source file list)\n",
        "annotation_list = sorted(glob.glob(dataset_directory + '/*.xml'))\n",
        "image_list = sorted(glob.glob(dataset_directory + '/*.jpg'))\n",
        "\n",
        "file_num = len(annotation_list)\n",
        "\n",
        "# インデックスシャッフル(shuffle)\n",
        "index_list = list(range(file_num - 1))\n",
        "random.shuffle(index_list)\n",
        "\n",
        "for count, index in enumerate(index_list):\n",
        "    if count < int(file_num * train_ratio):\n",
        "        # 学習用データ(Training Data)\n",
        "        shutil.copy2(annotation_list[index], train_directory)\n",
        "        shutil.copy2(image_list[index], train_directory)\n",
        "    else:\n",
        "        # 検証用データ(Validation Data)\n",
        "        shutil.copy2(annotation_list[index], validation_directory)\n",
        "        shutil.copy2(image_list[index], validation_directory)"
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ACKapHgx_d4Q"
      },
      "source": [
        "# Pascal VOC形式 を MS COCO形式へ変換(Convert Pascal VOC format to MS COCO format)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qKGpaUik_c9m"
      },
      "source": [
        "!git clone https://github.com/Kazuhito00/convert_voc_to_coco.git"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "t3xTEz30_kYp"
      },
      "source": [
        "!python convert_voc_to_coco/convert_voc_to_coco.py \\\n",
        "    train train/train_annotations.json \\\n",
        "    --start_image_id=0\n",
        "!python convert_voc_to_coco/convert_voc_to_coco.py \\\n",
        "    validation validation/validation_annotations.json \\\n",
        "    --start_image_id=10000000"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wJ9ytPB90pJP"
      },
      "source": [
        "# 学習データディレクトリ準備(Training data directory preparation)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IccyvWRpDZGL"
      },
      "source": [
        "!mkdir dataset\n",
        "!mkdir dataset/images\n",
        "!mkdir dataset/images/train2017\n",
        "!mkdir dataset/images/val2017\n",
        "!mkdir dataset/annotations\n",
        "\n",
        "!cp -rf train/*.jpg dataset/images/train2017\n",
        "!cp -rf validation/*.jpg dataset/images/val2017\n",
        "!cp -rf train/train_annotations.json dataset/annotations\n",
        "!cp -rf validation/validation_annotations.json dataset/annotations"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CnUirebA1a__"
      },
      "source": [
        "# コンフィグコピー\n",
        "<!--\n",
        "![image](https://user-images.githubusercontent.com/37477845/135283504-254ea817-345e-4665-828a-4c6034645ed1.png)\n",
        "-->\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gzlWZMuSPUly"
      },
      "source": [
        "if use_sample_image:\n",
        "    !cp /content/YOLOX-Colaboratory-Training-Sample/03.config/nano.py /content/YOLOX"
      ],
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vVvBXq4e2ydb"
      },
      "source": [
        "# モデル訓練"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ykzClTsh1ZDA"
      },
      "source": [
        "%cd /content/YOLOX/\n",
        "!wget https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_nano.pth"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!cp tools/train.py ./"
      ],
      "metadata": {
        "id": "HQGKhyNia21V"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qZDXBpaY22lN"
      },
      "source": [
        "!python train.py \\\n",
        "    -f nano.py \\\n",
        "    -d 1 \\\n",
        "    -b 16 \\\n",
        "    --fp16 \\\n",
        "    -o \\\n",
        "    -c yolox_nano.pth"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mmlvTpsheyQm"
      },
      "source": [
        "# 推論テスト(Inference test)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!cp tools/demo.py ./"
      ],
      "metadata": {
        "id": "69XkFyYwfu-L"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "w3g0ZRUwMP8k"
      },
      "source": [
        "TEST_IMAGE_PATH = \"/content/YOLOX-Colaboratory-Training-Sample/01.image/000050.jpg\"\n",
        "MODEL_PATH = \"/content/YOLOX/YOLOX_outputs/nano/best_ckpt.pth\"\n",
        "\n",
        "!python demo.py image \\\n",
        "    -f nano.py \\\n",
        "    -c {MODEL_PATH} \\\n",
        "    --path {TEST_IMAGE_PATH} \\\n",
        "    --conf 0.25 \\\n",
        "    --nms 0.45 \\\n",
        "    --tsize 640 \\\n",
        "    --save_result \\\n",
        "    --device gpu"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XL-oSdFrg0Pb"
      },
      "source": [
        "# from PIL import Image\n",
        "\n",
        "# OUTPUT_IMAGE_PATH = \"/content/YOLOX/YOLOX_outputs/nano/vis_res/2021_09_29_17_46_56/000050.jpg\" \n",
        "# Image.open(OUTPUT_IMAGE_PATH)"
      ],
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XT2WBgo7jAvR"
      },
      "source": [
        "# ONNX出力(Export ONNX Model)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!cp tools/export_onnx.py ./"
      ],
      "metadata": {
        "id": "MlmKqY6RfykM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vHpT0bQBhHzt"
      },
      "source": [
        "!python export_onnx.py \\\n",
        "    --output-name yolox_nano.onnx \\\n",
        "    -n yolox-nano \\\n",
        "    -f nano.py \\\n",
        "    -c {MODEL_PATH}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!cp demo/ONNXRuntime/onnx_inference.py ./"
      ],
      "metadata": {
        "id": "PYea1i3YgZE_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "q501MZh_jkIv"
      },
      "source": [
        "!python onnx_inference.py \\\n",
        "    -m yolox_nano.onnx \\\n",
        "    -i {TEST_IMAGE_PATH} \\\n",
        "    -o ./ \\\n",
        "    -s 0.3 \\\n",
        "    --input_shape 416,416"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uwcQysS_j_yp"
      },
      "source": [
        "from PIL import Image\n",
        "\n",
        "OUTPUT_IMAGE_PATH = \"000050.jpg\" \n",
        "Image.open(OUTPUT_IMAGE_PATH)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pxuKKcgj6WND"
      },
      "source": [
        "# ONNX -> TensorFlow 変換"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PKETcM79q3XD"
      },
      "source": [
        "!pip install onnx-tf"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EpiNXqOA6aoU"
      },
      "source": [
        "!onnx-tf convert \\\n",
        "    -i yolox_nano.onnx \\\n",
        "    -o yolox_nano_pb"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lV7E1wlX6cli"
      },
      "source": [
        "# TensorFlow -> TensorFlow-Lite 変換"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-Aln_cLaN5G0"
      },
      "source": [
        "!pip install tf-nightly"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "H1zYkqM26alJ"
      },
      "source": [
        "import tensorflow as tf"
      ],
      "execution_count": 26,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GI_X2zBhOSmw"
      },
      "source": [
        "%cd /content/YOLOX"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bdV9QUme-cD0"
      },
      "source": [
        "# ダイナミックレンジ量子化\n",
        "converter = tf.lite.TFLiteConverter.from_saved_model('yolox_nano_pb')\n",
        "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
        "tflite_quantized_model = converter.convert()\n",
        "\n",
        "open('yolox_nano_dynamic_range_quantize.tflite', 'wb').write(tflite_quantized_model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3wgIUaj06aiV"
      },
      "source": [
        "# 半精度浮動小数点数の量子化\n",
        "converter = tf.lite.TFLiteConverter.from_saved_model('yolox_nano_pb')\n",
        "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
        "converter.target_spec.supported_types = [tf.float16]\n",
        "tflite_quantized_model = converter.convert()\n",
        "\n",
        "open('yolox_nano_float16_quantize.tflite', 'wb').write(tflite_quantized_model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZjmUxifw6aew"
      },
      "source": [
        "# 完全整数量子化\n",
        "import glob\n",
        "import numpy as np\n",
        "from PIL import Image\n",
        "\n",
        "test_image_pathlist = glob.glob('/content/YOLOX-Colaboratory-Training-Sample/01.image/*.jpg')\n",
        "test_image_pathlist = test_image_pathlist[:100]\n",
        "print(len(test_image_pathlist))\n",
        "\n",
        "def representative_dataset():\n",
        "    for test_image_path in test_image_pathlist:\n",
        "        image = np.array(Image.open(test_image_path))\n",
        "        image = image.astype('float32')\n",
        "        image = tf.image.resize(image, (416, 416))\n",
        "        image = image - 127.5\n",
        "        image = image * 0.007843\n",
        "        image = tf.transpose(image, perm=[2, 0, 1])\n",
        "        image = tf.reshape(image, [1, 3, 416, 416])\n",
        "        yield [image]\n",
        "\n",
        "converter = tf.lite.TFLiteConverter.from_saved_model('yolox_nano_pb')\n",
        "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
        "converter.representative_dataset = representative_dataset\n",
        "tflite_quantized_model = converter.convert()\n",
        "\n",
        "open('yolox_nano_int8_quantize.tflite', 'wb').write(tflite_quantized_model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "y9cTtsmA6aaY"
      },
      "source": [
        "# 完全整数量子化(入力含む)\n",
        "import glob\n",
        "import numpy as np\n",
        "from PIL import Image\n",
        "\n",
        "test_image_pathlist = glob.glob('/content/YOLOX-Colaboratory-Training-Sample/01.image/*.jpg')\n",
        "test_image_pathlist = test_image_pathlist[:100]\n",
        "print(len(test_image_pathlist))\n",
        "\n",
        "def representative_dataset():\n",
        "    for test_image_path in test_image_pathlist:\n",
        "        image = np.array(Image.open(test_image_path))\n",
        "        image = image.astype('float32')\n",
        "        image = tf.image.resize(image, (416, 416))\n",
        "        image = image - 127.5\n",
        "        image = image * 0.007843\n",
        "        image = tf.transpose(image, perm=[2, 0, 1])\n",
        "        image = tf.reshape(image, [1, 3, 416, 416])\n",
        "        yield [image]\n",
        "\n",
        "converter = tf.lite.TFLiteConverter.from_saved_model('yolox_nano_pb')\n",
        "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
        "converter.representative_dataset = representative_dataset\n",
        "converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]\n",
        "converter.inference_input_type = tf.int8  # or tf.uint8\n",
        "converter.inference_output_type = tf.int8  # or tf.uint8\n",
        "tflite_quantized_model = converter.convert()\n",
        "\n",
        "open('yolox_nano_only_int8_quantize.tflite', 'wb').write(tflite_quantized_model)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}