11_cnn_pytorch/07_single_shot_multibox_detector.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "07_single_shot_multibox_detector",
      "provenance": [],
      "collapsed_sections": [
        "eCpsjyUN_npH",
        "_hT3k18MZwlH"
      ],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/machine-perception-robotics-group/MPRGDeepLearningLectureNotebook/blob/master/11_cnn_pytorch/07_single_shot_multibox_detector.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LSdX7WVvzGVN"
      },
      "source": [
        "# Single Shot Mutibox Detector (SSD)\n",
        "\n",
        "\n",
        "---\n",
        "\n",
        "\n",
        "物体検出タスクで主流となっているSSDについて説明します．"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4i6TTO8qBFiB"
      },
      "source": [
        "## 概要\n",
        "SSDは，物体候補領域の検出とクラス分類を単一のネットワークで行うモデルです．SSDのネットワーク構造は，複数の畳み込み層から物体矩形と物体クラスのスコアを出力させることで，スケール変化に頑健な物体検出ができます．6つの予測レイヤでは，アスペクト比が固定されたデフォルトボックスと呼ばれる矩形を使って複数の物体候補領域を検出します．複数のアスペクト比のデフォルトボックスを使うことで，様々な形状の物体に対応可能です．予測レイヤで得られるボックスのオフセットを表す特徴マップ(Localization maps)を用いて，すべてのデフォルトボックスをより物体に近い形状に回帰します．また，同様に予測レイヤで得られる物体のクラス尤度を表す特徴マップ(Confidence maps)を用いて，物体のクラスを識別します．\n",
        "\n",
        "\n",
        "<img src=\"https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2654719/52d5bb7b-a6c8-9c0a-77e0-3b0b431a0a90.png\" width = 100%>\n",
        "\n",
        "<img src=\"https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2654719/d16141e2-2994-d6c8-672b-bb196d187997.png\" width = 40%>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eCpsjyUN_npH"
      },
      "source": [
        "### Google Colaboratoryの設定確認・変更\n",
        "本チュートリアルではPyTorchを利用してニューラルネットワークの実装を確認，学習および評価を行います．\n",
        "**GPUを用いて処理を行うために，上部のメニューバーの「ランタイム」→「ランタイムのタイプを変更」からハードウェアアクセラレータをGPUにしてください．**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EdfN8rBnCbzq"
      },
      "source": [
        "# データセットの用意\n",
        "まず，学習を行うデータセットをダウンロードします．ここでは，物体検出タスクの性能評価に用いられるARC2017 RGB-Dデータセットを使用します．ARC2017 RGB-Dデータセットは2017年にAmazon Robotics Challengeにて使用された全40クラスを収録する物体検出のデータセットです．\n",
        "ARC2017 RGB-Dデータセットの詳細については以下サイトを参照してください．\n",
        "http://mprg.jp/research/arc_dataset_2017_j"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# データのダウンロード\n",
        "import gdown\n",
        "gdown.download('https://drive.google.com/uc?id=1oJs_uRkCG2eWBKKY7ekVJp2vlCZmTF1h', 'ARCdataset_png.zip', quiet=True)\n",
        "!unzip -q ARCdataset_png.zip"
      ],
      "metadata": {
        "id": "hh5v0wMUz6J8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VYXjCs303oGH"
      },
      "source": [
        "今回は，画像サイズは半分，枚数を削減したオリジナルのデータセットを使用します．\\\n",
        "40のオブジェクトと背景の全41クラスがラベル付けされています．\n",
        "\n",
        "<img src=\"https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2654719/6d9c3bc1-01d9-d433-a7fe-1b1df9aff2ce.png\" width = 70%>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UcdSHZnBDp3j"
      },
      "source": [
        "データセットの内容は以下のようになっています．\\\n",
        "データ構成: RGBの画像(640x480, .png)とバウンディングボックスデータ(.txt)\\\n",
        "学習データ -> 200枚 \\\n",
        "評価データ -> 50枚\n",
        "\n",
        "<img src=\"https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2654719/578e14bd-5514-3085-c5bb-376fd897d8d0.png\" width = 50%><img src=\"https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2654719/19783182-a96f-c165-4f10-4d6d36dc1e2c.png\" width = 50%>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gvdvyk56KGLd"
      },
      "source": [
        "ダウンロードしたデータの中身を確認します．\n",
        "フォルダARCdataset_pngには以下のフォルダ・ファイルが内包されています． \\\\\n",
        "- train ← 学習用データ\n",
        "- test ← 評価用データ\n",
        "- SSD_pretrained.pth ← 学習済みモデル \\\\"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uXiAu2RKExAG"
      },
      "source": [
        "作業ディレクトリを先程ダウンロードしたフォルダに移動します．"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# 作業ディレクトリの移動\n",
        "%cd ARCdataset_png\n",
        "!ls"
      ],
      "metadata": {
        "id": "hdjZMVvOz4Yq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L7DFz0Hv8Sal"
      },
      "source": [
        "## モジュールのインポート\n",
        "プログラムを実行するために必要なモジュールをインポートします．\n",
        "本実験では，facebookが開発した機械学習ライブラリのPyTorchを使用します．"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import datetime\n",
        "from glob import glob\n",
        "import os\n",
        "import sys\n",
        "import time\n",
        "from tqdm import tqdm\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import math\n",
        "import random\n",
        "import cv2 as cv\n",
        "import numpy as np\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torch.nn.functional as F\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "import torchsummary"
      ],
      "metadata": {
        "id": "-brs8jRxYIoC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K16YK_yYOMQN"
      },
      "source": [
        "実験はGPUを使用するので，colaboratory上で使用可能な状態か確認します．"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# GPUの確認\n",
        "use_cuda = torch.cuda.is_available()\n",
        "print('Use CUDA:', use_cuda)"
      ],
      "metadata": {
        "id": "V3LukARaz0D4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "シード値の固定を行います．"
      ],
      "metadata": {
        "id": "59W8URNwlb0f"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "seed = 123\n",
        "\n",
        "# Python random\n",
        "random.seed(seed)\n",
        "# Numpy\n",
        "np.random.seed(seed)\n",
        "# Pytorch\n",
        "torch.manual_seed(seed)\n",
        "torch.cuda.manual_seed(seed)\n",
        "torch.backends.cudnn.deterministic = True\n",
        "torch.use_deterministic_algorithms = True"
      ],
      "metadata": {
        "id": "QqTY7v7mliC7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mcEP0IK-9rvW"
      },
      "source": [
        "最初に，評価する画像が入ったフォルダやネットワークのハイパーパラメータを定義します．後に作成するネットワーク定義やテストプログラムからも，値を参照できるようにしておきます．"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class common:\n",
        "  def __init__(self):\n",
        "    # 学習画像のディレクトリ\n",
        "    self.images_dir = \".\"\n",
        "\n",
        "    # Data Augmentationにより作成されたデータのリストパス\n",
        "    self.augmented_img_list = \"./train/augimg_name_list.txt\"\n",
        "\n",
        "    # SSDの入力画像サイズ\n",
        "    self.insize = 300\n",
        "\n",
        "    # 識別のクラス数 (背景込み)\n",
        "    self.num_of_classes = 41\n",
        "\n",
        "    # Bounding boxのオフセットベクトルの次元数\n",
        "    self.num_of_offset_dims = 4\n",
        "\n",
        "    # Bounding boxのオフセットとクラスを推定する畳み込み層\n",
        "    self.mbox_source_layers = ['conv4_3', 'conv7', 'conv8_2', 'conv9_2', 'conv10_2', 'conv11_2']\n",
        "\n",
        "    # Default boxの最小・最大比率 (in percent %)\n",
        "    self.min_ratio = 20\n",
        "    self.max_ratio = 90\n",
        "\n",
        "    # 各階層における特徴マップの入力画像上のステップ幅\n",
        "    self.steps = [8, 16, 32, 64, 100, 300]\n",
        "\n",
        "    # 各階層のdefault boxの数\n",
        "    self.num_boxes = [4, 6, 6, 6, 4, 4]\n",
        "\n",
        "    # 各階層の特徴マップのDefault boxのアスペクト比\n",
        "    self.aspect_ratios = [[2], [2, 3], [2, 3], [2, 3], [2], [2]]\n",
        "\n",
        "    # 各階層の特徴マップの解像度\n",
        "    self.map_sizes = [38, 19, 10, 5, 3, 1]\n",
        "\n",
        "    self.loc_var = 0.01\n",
        "\n",
        "    # Data Augmentationで増加させる倍数\n",
        "    self.augmentation_factor = 2\n",
        "\n",
        "    # Data Augmentationのパラメータ\n",
        "    self.jitter = 0.3      # Default : 0.2\n",
        "    self.saturation = 1.5  # Default : 1.5\n",
        "    self.exposure = 1.5    # Default : 1.5\n",
        "    self.hue = 0.03        # Default : 0.03\n",
        "\n",
        "    # choose border type\n",
        "    # BORDER_REPLICATE, BORDER_REFLECT, BORDER_REFLECT_101, BORDER_WRAP, BORDER_CONSTANT\n",
        "    self.border_type = cv.BORDER_CONSTANT\n",
        "\n",
        "    self.border_val = (127, 127, 127)\n",
        "\n",
        "    # ARCクラスラベル (クラス名にはスペース(空白)は禁止)\n",
        "    self.arc_labels = [\n",
        "          \"Background\",             #0\n",
        "          \"Binder\",                 #1\n",
        "          \"Balloons\",               #2\n",
        "          \"Baby_Wipes\",             #3\n",
        "          \"Toilet_Brush\",           #4\n",
        "          \"Toothbrushes\",           #5\n",
        "          \"Crayons\",                #6\n",
        "          \"Salts\",                  #7\n",
        "          \"DVD\",                    #8\n",
        "          \"Glue_Sticks\",            #9\n",
        "          \"Eraser\",                 #10\n",
        "          \"Scissors\",               #11\n",
        "          \"Green_Book\",             #12\n",
        "          \"Socks\",                  #13\n",
        "          \"Irish_Spring\",           #14\n",
        "          \"Paper_Tape\",             #15\n",
        "          \"Touch_Tissues\",          #16\n",
        "          \"Knit_Gloves\",            #17\n",
        "          \"Laugh_Out_Loud_Jokes\",   #18\n",
        "          \"Pencil_Cup\",             #19\n",
        "          \"Mini_Marbles\",           #20\n",
        "          \"Neoprene_Weight\",        #21\n",
        "          \"Wine_Glasses\",           #22\n",
        "          \"Water_Bottle\",           #23\n",
        "          \"Reynolds_Pie\",           #24\n",
        "          \"Reynolds_Wrap\",          #25\n",
        "          \"Robots_Everywhere\",      #26\n",
        "          \"Duct_Tape\",              #27\n",
        "          \"Sponges\",                #28\n",
        "          \"Speed_Stick\",            #29\n",
        "          \"Index_Cards\",            #30\n",
        "          \"Ice_Cube_Tray\",          #31\n",
        "          \"Table_Cover\",            #32\n",
        "          \"Measuring_Spoons\",       #33\n",
        "          \"Bath_Sponge\",            #34\n",
        "          \"Pencils\",                #35\n",
        "          \"Mousetraps\",             #36\n",
        "          \"Face_Cloth\",             #37\n",
        "          \"Tennis_Balls\",           #38\n",
        "          \"Spray_Bottle\",           #39\n",
        "          \"Flashlights\"]            #40\n",
        "\n",
        "    # アイテムIDリスト\n",
        "    self.itemIDList = [\n",
        "                  \"0 BG\",\n",
        "                  \"1\",\n",
        "                  \"2\",\n",
        "                  \"3\",\n",
        "                  \"4\",\n",
        "                  \"5\",\n",
        "                  \"6\",\n",
        "                  \"7\",\n",
        "                  \"8\",\n",
        "                  \"9\",\n",
        "                  \"10\",\n",
        "                  \"11\",\n",
        "                  \"12\",\n",
        "                  \"13\",\n",
        "                  \"14\",\n",
        "                  \"15\",\n",
        "                  \"16\",\n",
        "                  \"17\",\n",
        "                  \"18\",\n",
        "                  \"19\",\n",
        "                  \"20\",\n",
        "                  \"21\",\n",
        "                  \"22\",\n",
        "                  \"23\",\n",
        "                  \"24\",\n",
        "                  \"25\",\n",
        "                  \"26\",\n",
        "                  \"27\",\n",
        "                  \"28\",\n",
        "                  \"29\",\n",
        "                  \"30\",\n",
        "                  \"31\",\n",
        "                  \"32\",\n",
        "                  \"33\",\n",
        "                  \"34\",\n",
        "                  \"35\",\n",
        "                  \"36\",\n",
        "                  \"37\",\n",
        "                  \"38\",\n",
        "                  \"39\",\n",
        "                  \"40\",\n",
        "                  \"41\"]\n",
        "\n",
        "    # クラスの色\n",
        "    self.arc_class_color = np.array([\n",
        "               [  0,   0,   0],\n",
        "               [ 85,   0,   0],\n",
        "               [170,   0,   0],\n",
        "               [255,   0,   0],\n",
        "               [  0,  85,   0],\n",
        "               [ 85,  85,   0],\n",
        "               [170,  85,   0],\n",
        "               [255,  85,   0],\n",
        "               [  0, 170,   0],\n",
        "               [ 85, 170,   0],\n",
        "               [170, 170,   0],\n",
        "               [255, 170,   0],\n",
        "               [  0, 255,   0],\n",
        "               [ 85, 255,   0],\n",
        "               [170, 255,   0],\n",
        "               [255, 255,   0],\n",
        "               [  0,   0,  85],\n",
        "               [ 85,   0,  85],\n",
        "               [170,   0,  85],\n",
        "               [255,   0,  85],\n",
        "               [  0,  85,  85],\n",
        "               [ 85,  85,  85],\n",
        "               [170,  85,  85],\n",
        "               [255,  85,  85],\n",
        "               [  0, 170,  85],\n",
        "               [ 85, 170,  85],\n",
        "               [170, 170,  85],\n",
        "               [255, 170,  85],\n",
        "               [  0, 255,  85],\n",
        "               [ 85, 255,  85],\n",
        "               [170, 255,  85],\n",
        "               [255, 255,  85],\n",
        "               [  0,   0, 170],\n",
        "               [ 85,   0, 170],\n",
        "               [170,   0, 170],\n",
        "               [255,   0, 170],\n",
        "               [  0,  85, 170],\n",
        "               [ 85,  85, 170],\n",
        "               [170,  85, 170],\n",
        "               [255,  85, 170],\n",
        "               [  0, 170, 170]])\n",
        "\n",
        "common_params = common()"
      ],
      "metadata": {
        "id": "VqXA6u1EYLer"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Cjhcl8TB-Onr"
      },
      "source": [
        "## データ拡張\n",
        "Data Augmentationの処理を記述します．trainフォルダ内にデータ拡張済みのBoundingBoxファイルが用意されており，データ拡張済みのファイルが読み込まれた際に画像に対してData Augmentation処理が実行されます．\\\n",
        "読み込まれたRGB画像は一度HSV形式に変換され，歪みを加えたRGB画像を新たに生成します．"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# BGR -> HSV\n",
        "def BGR_2_HSV_(img):\n",
        "\n",
        "    h = np.zeros((img.shape[0], img.shape[1]), np.float32)\n",
        "    s = np.zeros((img.shape[0], img.shape[1]), np.float32)\n",
        "    v = np.zeros((img.shape[0], img.shape[1]), np.float32)\n",
        "\n",
        "    img_r = img.copy()\n",
        "    b = img[:, :, 0]\n",
        "    g = img[:, :, 1]\n",
        "    r = img[:, :, 2]\n",
        "    max_v = cv.max(cv.max(b, g), r)\n",
        "    min_v = cv.min(cv.min(b, g), r)\n",
        "    delta = max_v - min_v\n",
        "    v = max_v\n",
        "    zero_m = (max_v == 0.)\n",
        "    zero_m = zero_m.astype(np.float32)\n",
        "    nonzeor_m = (max_v != 0.)\n",
        "    nonzeor_m = nonzeor_m.astype(np.float32)\n",
        "\n",
        "    exp_m = zero_m * 10e-8\n",
        "\n",
        "    s = delta / (max_v + exp_m)\n",
        "    s = s * nonzeor_m\n",
        "\n",
        "    rmax = (r == max_v)\n",
        "    rmax = rmax.astype(np.float32)\n",
        "\n",
        "    gmax = (g == max_v)\n",
        "    gr = (g != r)\n",
        "    gmax = gmax.astype(np.float32)\n",
        "    gr = gr.astype(np.float32)\n",
        "    gmax = gmax * gr\n",
        "\n",
        "    bmax = (b == max_v)\n",
        "    br = (b != r)\n",
        "    bg = (b != g)\n",
        "    bmax = bmax.astype(np.float32)\n",
        "    br = br.astype(np.float32)\n",
        "    bg = bg.astype(np.float32)\n",
        "    bmax = bmax * br * bg\n",
        "\n",
        "    h += ((g - b) / (delta + 10e-8)) * rmax\n",
        "    h += (((b - r) / (delta + 10e-8)) + 2.) * gmax\n",
        "    h += (((r - g) / (delta + 10e-8)) + 4.) * bmax\n",
        "\n",
        "    h = h * (np.pi / 3.)\n",
        "\n",
        "    neg_m = (h < 0.0)\n",
        "    neg_m = neg_m.astype(np.float32)\n",
        "    h += neg_m * (2. * np.pi)\n",
        "\n",
        "    h = h / (2. * np.pi)\n",
        "\n",
        "    img_r[:, :, 0] = h\n",
        "    img_r[:, :, 1] = s\n",
        "    img_r[:, :, 2] = v\n",
        "\n",
        "    return img_r\n",
        "\n",
        "# HSV -> BGR\n",
        "def HSV_2_BGR_(img):\n",
        "\n",
        "    r = np.zeros((img.shape[0], img.shape[1]), np.float32)\n",
        "    g = np.zeros((img.shape[0], img.shape[1]), np.float32)\n",
        "    b = np.zeros((img.shape[0], img.shape[1]), np.float32)\n",
        "\n",
        "    img_r = img.copy()\n",
        "    h = img[:, :, 0]\n",
        "    s = img[:, :, 1]\n",
        "    v = img[:, :, 2]\n",
        "\n",
        "    h = h * (2. * np.pi)\n",
        "\n",
        "    zero_m = (s == 0.)\n",
        "    zero_m = zero_m.astype(np.float32)\n",
        "    nonzeor_m = (s != 0.)\n",
        "    nonzeor_m = nonzeor_m.astype(np.float32)\n",
        "\n",
        "    idx = np.floor(h)\n",
        "    idx = idx.astype(np.int16)\n",
        "\n",
        "    f = h - idx\n",
        "    p = v * (1. - s)\n",
        "    q = v * (1. - s * ((3. / np.pi) * f))\n",
        "    t = v * (1. - s * (1. - ((3. / np.pi) * f)))\n",
        "\n",
        "    idx0 = (idx == 0)\n",
        "    idx0 = idx0.astype(np.float32)\n",
        "\n",
        "    idx1 = (idx == 1)\n",
        "    idx1 = idx1.astype(np.float32)\n",
        "\n",
        "    idx2 = (idx == 2)\n",
        "    idx2 = idx2.astype(np.float32)\n",
        "\n",
        "    idx3 = (idx == 3)\n",
        "    idx3 = idx3.astype(np.float32)\n",
        "\n",
        "    idx4 = (idx == 4)\n",
        "    idx4 = idx4.astype(np.float32)\n",
        "\n",
        "    idxE = idx0 + idx1 + idx2 + idx3 + idx4\n",
        "    idxE = (idxE == 0)\n",
        "    idxE = idxE.astype(np.float32)\n",
        "\n",
        "    r += v * idx0\n",
        "    g += t * idx0\n",
        "    b += p * idx0\n",
        "\n",
        "    r += q * idx1\n",
        "    g += v * idx1\n",
        "    b += p * idx1\n",
        "\n",
        "    r += p * idx2\n",
        "    g += v * idx2\n",
        "    b += t * idx2\n",
        "\n",
        "    r += p * idx3\n",
        "    g += q * idx3\n",
        "    b += v * idx3\n",
        "\n",
        "    r += t * idx4\n",
        "    g += p * idx4\n",
        "    b += v * idx4\n",
        "\n",
        "    r += v * idxE\n",
        "    g += p * idxE\n",
        "    b += q * idxE\n",
        "\n",
        "    r = r * nonzeor_m\n",
        "    g = g * nonzeor_m\n",
        "    b = b * nonzeor_m\n",
        "\n",
        "    r += v * zero_m\n",
        "    g += v * zero_m\n",
        "    b += v * zero_m\n",
        "\n",
        "    img_r[:, :, 0] = b\n",
        "    img_r[:, :, 1] = g\n",
        "    img_r[:, :, 2] = r\n",
        "\n",
        "    return img_r\n",
        "\n",
        "# 歪みをくわえた画像を生成\n",
        "def distortImage(img, dhue, dsat, dexp):\n",
        "\n",
        "    img = BGR_2_HSV_(img)\n",
        "\n",
        "    n = img.shape[0] * img.shape[1]\n",
        "\n",
        "    img[:, :, 0] = img[:, :, 0] + dhue\n",
        "    img[:, :, 1] = img[:, :, 1] * dsat\n",
        "    img[:, :, 2] = img[:, :, 2] * dexp\n",
        "\n",
        "    m = img[:, :, 0] > 1.\n",
        "    m = m.astype(np.float32)\n",
        "    p = img[:, :, 0] < 0.\n",
        "    p = p.astype(np.float32)\n",
        "\n",
        "    img[:, :, 0] = img[:, :, 0] - m\n",
        "    img[:, :, 0] = img[:, :, 0] + p\n",
        "\n",
        "    img = HSV_2_BGR_(img)\n",
        "\n",
        "    img = np.minimum(np.maximum(img, 0.), 1.)\n",
        "\n",
        "    return img\n",
        "\n",
        "\n",
        "# Data augmentation 処理\n",
        "def trainAugmentation(img, border_pixels, crop_param, hsv_param, flip_type):\n",
        "\n",
        "    hsv_srand = 0\n",
        "\n",
        "    if hsv_srand == 1:\n",
        "        dhue = hsv_param[0]\n",
        "        dsat = hsv_param[1]\n",
        "        dexp = hsv_param[2]\n",
        "    else:\n",
        "        dhue = random.uniform(-1. * common_params.hue, common_params.hue)\n",
        "        dsat = random.uniform(1., common_params.saturation)\n",
        "        dexp = random.uniform(1., common_params.exposure)\n",
        "        if random.randint(0, 1) == 0:\n",
        "            dsat = 1. / dsat\n",
        "        if random.randint(0, 1) == 0:\n",
        "            dexp = 1. / dexp\n",
        "\n",
        "    if random.randint(0, 4) >= 1:\n",
        "        img = img.astype(np.float32)\n",
        "        img /= 255.\n",
        "        img = distortImage(img, dhue, dsat, dexp)\n",
        "        img *= 255.\n",
        "        img = img.astype(np.uint8)\n",
        "\n",
        "    if common_params.border_type == cv.BORDER_CONSTANT:\n",
        "        cropped_img = cv.copyMakeBorder(img, border_pixels[1], border_pixels[3], border_pixels[0], border_pixels[2], common_params.border_type, value = common_params.border_val)\n",
        "    else:\n",
        "        cropped_img = cv.copyMakeBorder(img, border_pixels[1], border_pixels[3], border_pixels[0], border_pixels[2], common_params.border_type)\n",
        "\n",
        "    cropped_img = cv.getRectSubPix(cropped_img, (int(crop_param[2]), int(crop_param[3])), (crop_param[0], crop_param[1]))\n",
        "    cropped_img = cv.resize(cropped_img, (common_params.insize, common_params.insize), interpolation = cv.INTER_CUBIC)\n",
        "\n",
        "    if flip_type >= 0:\n",
        "        cropped_img = cv.flip(cropped_img, flip_type)\n",
        "\n",
        "    return cropped_img"
      ],
      "metadata": {
        "id": "sA-0SgKeYQbe"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E4iklfyAG5pf"
      },
      "source": [
        "## データローダの作成\n",
        "データセットの読み込み方を定義したデータローダを作成します．\n",
        "オリジナルのデータセットを使用する場合に必要です．\\\n",
        "今回は`transforms.Compose`のデータ拡張機能は使用せずに，MYDatasetクラスを`DataLoader`に与えてデータを読み込みます． \\\n",
        "バッチサイズは5，データセット読み込みに割り当てるスレッド数(num_workers)は2とします．"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# バッチサイズの設定\n",
        "batch_size = 5\n",
        "\n",
        "class MYDataset(Dataset):\n",
        "    def __init__(self, input_list_path, confing_image):\n",
        "        f = open(input_list_path, 'r')\n",
        "        input_list = []\n",
        "        for line in f:\n",
        "            ln = line.split('\\n')\n",
        "            input_list.append(ln[0])\n",
        "        input_list.sort()\n",
        "        f.close()\n",
        "\n",
        "        self.input_list = np.array(input_list)\n",
        "        self.confing_image = confing_image\n",
        "        print(\"Training images : \", len(self.input_list))\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.input_list)\n",
        "\n",
        "    def __getitem__(self, i):\n",
        "        input_name = self.input_list[i]\n",
        "\n",
        "        # Detection teach\n",
        "        aug_p = open(common_params.images_dir + \"/train/img_aug_param/\" + input_name + \".txt\", 'r')\n",
        "\n",
        "        in_line = aug_p.readline()  #aug_pを1行ずつ読み込み\n",
        "        opath = in_line.split(' \\n')    #改行で区切る\n",
        "        original_img_path = str(opath[0])\n",
        "\n",
        "        in_line = aug_p.readline()\n",
        "        augmentation = int(in_line)\n",
        "\n",
        "        in_line = aug_p.readline()\n",
        "        part = in_line.split(' ')\n",
        "        border_pixels = [int(part[0]), int(part[1]), int(part[2]), int(part[3])]\n",
        "\n",
        "        in_line = aug_p.readline()\n",
        "        part = in_line.split(' ')\n",
        "        crop_param = [float(part[0]), float(part[1]), float(part[2]), float(part[3])]\n",
        "\n",
        "        in_line = aug_p.readline()\n",
        "        part = in_line.split(' ')\n",
        "        hsv_param = [float(part[0]), float(part[1]), float(part[2])]\n",
        "\n",
        "        in_line = aug_p.readline()\n",
        "        flip_type = int(in_line)\n",
        "\n",
        "        # Input image\n",
        "        color_img = cv.imread(common_params.images_dir + \"/train/rgb/\" + original_img_path + \".png\", cv.IMREAD_COLOR)\n",
        "\n",
        "        # 画像をSSDの入力サイズにリサイズ\n",
        "        input_img = cv.resize(color_img, (common_params.insize, common_params.insize), interpolation = cv.INTER_CUBIC)  #バイキュビック補間\n",
        "\n",
        "        # Data Augmentationにより作成されたBoundingBox情報か判定\n",
        "        if augmentation == 1:\n",
        "            input_img = trainAugmentation(input_img, border_pixels, crop_param, hsv_param, flip_type)   #data augmentation\n",
        "\n",
        "        if self.confing_image:\n",
        "            conf_img = input_img.copy()\n",
        "\n",
        "        # 画像データをfloatに変換\n",
        "        input_img = input_img.astype(np.float32)\n",
        "\n",
        "        # 画像の平均値を引く\n",
        "        input_img -= np.array([103.939, 116.779, 123.68])\n",
        "\n",
        "        # 画像の次元を(高さ，幅，チャンネル数)から(チャンネル数, 高さ，幅)へ転置\n",
        "        input_img = input_img.transpose(2, 0, 1)\n",
        "\n",
        "        gt_boxes = []\n",
        "        df_boxes = []\n",
        "        indices = []\n",
        "        classes = []\n",
        "        idx_tmp = []\n",
        "\n",
        "        # positiveサンプルの読み込み\n",
        "        pos_num = 0\n",
        "        f = open(common_params.images_dir + \"/train/positives/\" + input_name + \".txt\", 'r')\n",
        "        for rw in f:\n",
        "            ln = rw.split(' ')\n",
        "            classes.append(int(ln[1]))\n",
        "            gt_boxes.append([float(ln[2]), float(ln[3]), float(ln[4]), float(ln[5])])\n",
        "            df_boxes.append([float(ln[6]), float(ln[7]), float(ln[8]), float(ln[9])])\n",
        "            indices.append([int(ln[10]), int(ln[11]), int(ln[12]), int(ln[13])])\n",
        "            pos_num += 1\n",
        "        f.close()\n",
        "\n",
        "        # hard negativeサンプルの読み込み (最大でpositiveサンプル数の3倍)\n",
        "        neg_num = 0\n",
        "        f = open(common_params.images_dir + \"/train/negatives/\" + input_name + \".txt\", 'r')\n",
        "        for rw in f:\n",
        "            ln = rw.split(' ')\n",
        "            classes.append(int(ln[1]))\n",
        "            gt_boxes.append([float(ln[2]), float(ln[3]), float(ln[4]), float(ln[5])])\n",
        "            df_boxes.append([float(ln[2]), float(ln[3]), float(ln[4]), float(ln[5])])\n",
        "            idx_tmp.append([int(ln[10]), int(ln[11]), int(ln[12]), int(ln[13])])\n",
        "            neg_num += 1\n",
        "        f.close()\n",
        "\n",
        "        hardneg_size = pos_num * 3 if neg_num > (pos_num * 3) else neg_num\n",
        "\n",
        "        perm = np.random.permutation(len(idx_tmp))\n",
        "\n",
        "        for hn in range(0, hardneg_size):\n",
        "            indices.append(idx_tmp[perm[hn]])\n",
        "\n",
        "        # padding処理 ???(random) -> 8732(dfbox max size)\n",
        "        max_boxes = 8732 * 4\n",
        "\n",
        "        if len(gt_boxes) == 0:\n",
        "            gt_boxes = np.ones((max_boxes, 4)) * -100\n",
        "        elif len(gt_boxes) != max_boxes:\n",
        "            gt_boxes = np.pad(np.array(gt_boxes), [(0,max_boxes-len(gt_boxes)), (0,0)], 'constant', constant_values=-100)\n",
        "\n",
        "        if len(df_boxes) == 0:\n",
        "            df_boxes = np.ones((max_boxes, 4)) * -100\n",
        "        elif len(df_boxes) != max_boxes:\n",
        "            df_boxes = np.pad(np.array(df_boxes), [(0,max_boxes-len(df_boxes)), (0,0)], 'constant', constant_values=-100)\n",
        "\n",
        "        if len(indices) == 0:\n",
        "            indices = np.ones((max_boxes, 4)) * -100\n",
        "        elif len(indices) != max_boxes:\n",
        "            indices = np.pad(np.array(indices), [(0,max_boxes-len(indices)), (0,0)], 'constant', constant_values=-100)\n",
        "\n",
        "        if len(classes) == 0:\n",
        "            classes = np.ones((max_boxes,)) * -100\n",
        "        elif len(classes) != max_boxes:\n",
        "            classes = np.pad(np.array(classes), (0,max_boxes-len(classes)), 'constant', constant_values=-100)\n",
        "\n",
        "        # list to tensor\n",
        "        input_img = torch.tensor(input_img, dtype=torch.float32).cpu()\n",
        "        gt_boxes = torch.tensor(gt_boxes, dtype=torch.float32).cpu()\n",
        "        df_boxes = torch.tensor(df_boxes, dtype=torch.float32).cpu()\n",
        "        indices = torch.tensor(indices, dtype=torch.int64).cpu()\n",
        "        classes = torch.tensor(classes, dtype=torch.int64).cpu()\n",
        "        conf_img = torch.tensor(conf_img).cpu()\n",
        "\n",
        "        return input_img, gt_boxes, df_boxes, indices, classes, conf_img\n",
        "\n",
        "# Dataset import\n",
        "train_dataset = MYDataset(common_params.augmented_img_list, True)\n",
        "train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2)"
      ],
      "metadata": {
        "id": "Tn5VTXb1zviQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t__FbTYANd5S"
      },
      "source": [
        "# ネットワークモデルの定義\n",
        "次に，SSDのネットワークを定義します．SSDは入力サイズが300x300のSSD300と512x512のSSD512があります．SSD512の方が入力サイズが大きいため，基本的には検出精度は向上しますが処理に時間がかかります．今回の実験では，入力サイズの小さいSSD300を用います．\\\n",
        "SSDはVGGをベースとしており，ネットワークの前段で小さい物体の特徴を捉え，後段で大きい物体の特徴を捉えます．\\\n",
        "SSD300では畳み込み層の(`Conv4_3,`,`Conv7`,`Conv8_2`,`Conv9_2`,`Conv10_2`,`Conv11_2`)からBounding boxのオフセットとクラスを推定します．\n",
        "また，`Conv4_3`の時のみL2正規化後(`L2Normalization`)の出力値を推定する層に用います．これは入力層に近い層のスケールが大きくなるのを防ぐためです．"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class L2Normalization(nn.Module):  # L2正規化\n",
        "    def __init__(self, channels, gamma=20):\n",
        "        super().__init__()\n",
        "        self.gamma = gamma\n",
        "        self.in_channels = channels\n",
        "        self.out_channels = channels\n",
        "        self.scales = nn.Parameter(torch.Tensor(self.in_channels))  # trainable\n",
        "        self.reset_parameters()\n",
        "\n",
        "    def reset_parameters(self):\n",
        "        nn.init.constant_(self.scales, self.gamma)  # initialized with gamma first\n",
        "\n",
        "    # Note that pytorch's dimension order is batch_size, channels, height, width\n",
        "    def forward(self, x):\n",
        "        # |x|_2\n",
        "        # normalize (x^)\n",
        "        x = F.normalize(x, p=2, dim=1)\n",
        "        return self.scales.unsqueeze(0).unsqueeze(2).unsqueeze(3) * x\n",
        "\n",
        "\n",
        "class SSDNet(nn.Module):  # SSDのネットワーク構築\n",
        "\n",
        "    def __init__(self):\n",
        "        super(SSDNet, self).__init__()\n",
        "        # feature layers\n",
        "        self.conv1_1 = nn.Conv2d(3, 64, kernel_size = 3, stride = 1, padding = 1)\n",
        "        self.relu1_1 = nn.ReLU(inplace=True)\n",
        "        self.conv1_2 = nn.Conv2d(64, 64, kernel_size = 3, stride = 1, padding = 1)\n",
        "        self.relu1_2 = nn.ReLU(inplace=True)\n",
        "\n",
        "        self.maxpool1 = nn.MaxPool2d(kernel_size = 2, stride = 2, padding = 0)\n",
        "\n",
        "        self.conv2_1 = nn.Conv2d(64, 128, kernel_size = 3, stride = 1, padding = 1)\n",
        "        self.relu2_1 = nn.ReLU(inplace=True)\n",
        "        self.conv2_2 = nn.Conv2d(128, 128, kernel_size = 3, stride = 1, padding = 1)\n",
        "        self.relu2_2 = nn.ReLU(inplace=True)\n",
        "\n",
        "        self.maxpool2 = nn.MaxPool2d(kernel_size = 2, stride = 2, padding = 0)\n",
        "\n",
        "        self.conv3_1 = nn.Conv2d(128, 256, kernel_size = 3, stride = 1, padding = 1)\n",
        "        self.relu3_1 = nn.ReLU(inplace=True)\n",
        "        self.conv3_2 = nn.Conv2d(256, 256, kernel_size = 3, stride = 1, padding = 1)\n",
        "        self.relu3_2 = nn.ReLU(inplace=True)\n",
        "        self.conv3_3 = nn.Conv2d(256, 256, kernel_size = 3, stride = 1, padding = 1)\n",
        "        self.relu3_3 = nn.ReLU(inplace=True)\n",
        "\n",
        "        self.maxpool3 = nn.MaxPool2d(kernel_size = 2, stride = 2, padding = 0, ceil_mode = True)\n",
        "\n",
        "        self.conv4_1 = nn.Conv2d(256, 512, kernel_size = 3, stride = 1, padding = 1)\n",
        "        self.relu4_1 = nn.ReLU(inplace=True)\n",
        "        self.conv4_2 = nn.Conv2d(512, 512, kernel_size = 3, stride = 1, padding = 1)\n",
        "        self.relu4_2 = nn.ReLU(inplace=True)\n",
        "        self.conv4_3 = nn.Conv2d(512, 512, kernel_size = 3, stride = 1, padding = 1)\n",
        "        self.relu4_3 = nn.ReLU(inplace=True)\n",
        "\n",
        "        self.l2norm4_3 = L2Normalization(channels=512, gamma=20)\n",
        "\n",
        "        self.maxpool4 = nn.MaxPool2d(kernel_size = 2, stride = 2, padding = 0)\n",
        "\n",
        "        self.conv5_1 = nn.Conv2d(512, 512, kernel_size = 3, stride = 1, padding = 1)\n",
        "        self.relu5_1 = nn.ReLU(inplace=True)\n",
        "        self.conv5_2 = nn.Conv2d(512, 512, kernel_size = 3, stride = 1, padding = 1)\n",
        "        self.relu5_2 = nn.ReLU(inplace=True)\n",
        "        self.conv5_3 = nn.Conv2d(512, 512, kernel_size = 3, stride = 1, padding = 1)\n",
        "        self.relu5_3 = nn.ReLU(inplace=True)\n",
        "\n",
        "        self.maxpool5 = nn.MaxPool2d(kernel_size = 3, stride = 1, padding = 1)\n",
        "\n",
        "        self.conv6 = nn.Conv2d(512, 1024, kernel_size = 3, stride = 1, padding = 6, dilation=6)\n",
        "        self.relu6 = nn.ReLU(inplace=True)\n",
        "        self.conv7 = nn.Conv2d(1024, 1024, kernel_size = 1, stride = 1, padding = 0)\n",
        "        self.relu7 = nn.ReLU(inplace=True)\n",
        "\n",
        "        self.conv8_1 = nn.Conv2d(1024, 256, kernel_size=1, stride=1, padding=0)\n",
        "        self.relu8_1 = nn.ReLU(inplace=True)\n",
        "        self.conv8_2 = nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1)\n",
        "        self.relu8_2 = nn.ReLU(inplace=True)\n",
        "\n",
        "        self.conv9_1 = nn.Conv2d(512, 128, kernel_size=1, stride=1, padding=0)\n",
        "        self.relu9_1 = nn.ReLU(inplace=True)\n",
        "        self.conv9_2 = nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1)\n",
        "        self.relu9_2 = nn.ReLU(inplace=True)\n",
        "\n",
        "        self.conv10_1 = nn.Conv2d(256, 128, kernel_size=1, stride=1, padding=0)\n",
        "        self.relu10_1 = nn.ReLU(inplace=True)\n",
        "        self.conv10_2 = nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=0)\n",
        "        self.relu10_2 = nn.ReLU(inplace=True)\n",
        "\n",
        "        self.conv11_1 = nn.Conv2d(256, 128, kernel_size=1, stride=1, padding=0)\n",
        "        self.relu11_1 = nn.ReLU(inplace=True)\n",
        "        self.conv11_2 = nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=0)\n",
        "        self.relu11_2 = nn.ReLU(inplace=True)\n",
        "\n",
        "        # localization layers\n",
        "        self.conv4_3_loc = nn.Conv2d(512, common_params.num_of_offset_dims * common_params.num_boxes[0], kernel_size=3, stride=1, padding=1)\n",
        "        self.conv7_loc = nn.Conv2d(1024, common_params.num_of_offset_dims * common_params.num_boxes[1], kernel_size=3, stride=1, padding=1)\n",
        "        self.conv8_2_loc = nn.Conv2d(512, common_params.num_of_offset_dims * common_params.num_boxes[2], kernel_size=3, stride=1, padding=1)\n",
        "        self.conv9_2_loc = nn.Conv2d(256, common_params.num_of_offset_dims * common_params.num_boxes[3], kernel_size=3, stride=1, padding=1)\n",
        "        self.conv10_2_loc = nn.Conv2d(256, common_params.num_of_offset_dims * common_params.num_boxes[4], kernel_size=3, stride=1, padding=1)\n",
        "        self.conv11_2_loc = nn.Conv2d(256, common_params.num_of_offset_dims * common_params.num_boxes[5], kernel_size=3, stride=1, padding=1)\n",
        "\n",
        "        # confidence layers\n",
        "        self.conv4_3_cls = nn.Conv2d(512, common_params.num_of_classes * common_params.num_boxes[0], kernel_size=3, stride=1, padding=1)\n",
        "        self.conv7_cls = nn.Conv2d(1024, common_params.num_of_classes * common_params.num_boxes[1], kernel_size=3, stride=1, padding=1)\n",
        "        self.conv8_2_cls = nn.Conv2d(512, common_params.num_of_classes * common_params.num_boxes[2], kernel_size=3, stride=1, padding=1)\n",
        "        self.conv9_2_cls = nn.Conv2d(256, common_params.num_of_classes * common_params.num_boxes[3], kernel_size=3, stride=1, padding=1)\n",
        "        self.conv10_2_cls = nn.Conv2d(256, common_params.num_of_classes * common_params.num_boxes[4], kernel_size=3, stride=1, padding=1)\n",
        "        self.conv11_2_cls = nn.Conv2d(256, common_params.num_of_classes * common_params.num_boxes[5], kernel_size=3, stride=1, padding=1)\n",
        "\n",
        "    def forward(self, x, train=False):\n",
        "\n",
        "        k1 = self.relu1_1(self.conv1_1(x))\n",
        "        k1 = self.relu1_2(self.conv1_2(k1))\n",
        "\n",
        "        k1 = self.maxpool1(k1)\n",
        "\n",
        "        k1 = self.relu2_1(self.conv2_1(k1))\n",
        "        k1 = self.relu2_2(self.conv2_2(k1))\n",
        "\n",
        "        k1 = self.maxpool2(k1)\n",
        "\n",
        "        k1 = self.relu3_1(self.conv3_1(k1))\n",
        "        k1 = self.relu3_2(self.conv3_2(k1))\n",
        "        k1 = self.relu3_3(self.conv3_3(k1))\n",
        "\n",
        "        k1 = self.maxpool3(k1)\n",
        "\n",
        "        k1 = self.relu4_1(self.conv4_1(k1))\n",
        "        k1 = self.relu4_2(self.conv4_2(k1))\n",
        "        k1 = self.relu4_3(self.conv4_3(k1))\n",
        "\n",
        "        l2_k1 = self.l2norm4_3(k1)\n",
        "        Loc1 = self.conv4_3_loc(l2_k1)  # Box_Estimator_1\n",
        "        Cls1 = self.conv4_3_cls(l2_k1)  # Class_Classifier_1\n",
        "\n",
        "        k2 = self.maxpool4(k1)\n",
        "\n",
        "        k2 = self.relu5_1(self.conv5_1(k2))\n",
        "        k2 = self.relu5_2(self.conv5_2(k2))\n",
        "        k2 = self.relu5_3(self.conv5_3(k2))\n",
        "\n",
        "        k2 = self.maxpool5(k2)\n",
        "\n",
        "        k2 = self.relu6(self.conv6(k2))\n",
        "        k2 = self.relu7(self.conv7(k2))\n",
        "\n",
        "        Loc2 = self.conv7_loc(k2)  # Box_Estimator_2\n",
        "        Cls2 = self.conv7_cls(k2)  # Class_Classifier_2\n",
        "\n",
        "        k3 = self.relu8_1(self.conv8_1(k2))\n",
        "        k3 = self.relu8_2(self.conv8_2(k3))\n",
        "\n",
        "        Loc3 = self.conv8_2_loc(k3)  # Box_Estimator_3\n",
        "        Cls3 = self.conv8_2_cls(k3)  # Class_Classifier_3\n",
        "\n",
        "        k4 = self.relu9_1(self.conv9_1(k3))\n",
        "        k4 = self.relu9_2(self.conv9_2(k4))\n",
        "\n",
        "        Loc4 = self.conv9_2_loc(k4)  # Box_Estimator_4\n",
        "        Cls4 = self.conv9_2_cls(k4)  # Class_Classifier_4\n",
        "\n",
        "        k5 = self.relu10_1(self.conv10_1(k4))\n",
        "        k5 = self.relu10_2(self.conv10_2(k5))\n",
        "\n",
        "        Loc5 = self.conv10_2_loc(k5)  # Box_Estimator_5\n",
        "        Cls5 = self.conv10_2_cls(k5)  # Class_Classifier_5\n",
        "\n",
        "        k6 = self.relu11_1(self.conv11_1(k5))\n",
        "        k6 = self.relu11_2(self.conv11_2(k6))\n",
        "\n",
        "        Loc6 = self.conv11_2_loc(k6)  # Box_Estimator_6\n",
        "        Cls6 = self.conv11_2_cls(k6)  # Class_Classifier_6\n",
        "\n",
        "        if train:\n",
        "            # (バッチ数,チャンネル数,高さ,幅) -> (バッチ数,高さ,幅,チャンネル数) メモリ上で要素順に並べ直す\n",
        "            Loc1 = Loc1.permute(0, 2, 3, 1).contiguous()\n",
        "            Cls1 = Cls1.permute(0, 2, 3, 1).contiguous()\n",
        "\n",
        "            Loc2 = Loc2.permute(0, 2, 3, 1).contiguous()\n",
        "            Cls2 = Cls2.permute(0, 2, 3, 1).contiguous()\n",
        "\n",
        "            Loc3 = Loc3.permute(0, 2, 3, 1).contiguous()\n",
        "            Cls3 = Cls3.permute(0, 2, 3, 1).contiguous()\n",
        "\n",
        "            Loc4 = Loc4.permute(0, 2, 3, 1).contiguous()\n",
        "            Cls4 = Cls4.permute(0, 2, 3, 1).contiguous()\n",
        "\n",
        "            Loc5 = Loc5.permute(0, 2, 3, 1).contiguous()\n",
        "            Cls5 = Cls5.permute(0, 2, 3, 1).contiguous()\n",
        "\n",
        "            Loc6 = Loc6.permute(0, 2, 3, 1).contiguous()\n",
        "            Cls6 = Cls6.permute(0, 2, 3, 1).contiguous()\n",
        "\n",
        "            Loc1 = Loc1.view(Loc1.data.shape[0] * Loc1.data.shape[1] * Loc1.data.shape[2] * common_params.num_boxes[0], int(Loc1.data.shape[3] / common_params.num_boxes[0]))\n",
        "            Cls1 = Cls1.view(Cls1.data.shape[0] * Cls1.data.shape[1] * Cls1.data.shape[2] * common_params.num_boxes[0], int(Cls1.data.shape[3] / common_params.num_boxes[0]))\n",
        "\n",
        "            Loc2 = Loc2.view(Loc2.data.shape[0] * Loc2.data.shape[1] * Loc2.data.shape[2] * common_params.num_boxes[1], int(Loc2.data.shape[3] / common_params.num_boxes[1]))\n",
        "            Cls2 = Cls2.view(Cls2.data.shape[0] * Cls2.data.shape[1] * Cls2.data.shape[2] * common_params.num_boxes[1], int(Cls2.data.shape[3] / common_params.num_boxes[1]))\n",
        "\n",
        "            Loc3 = Loc3.view(Loc3.data.shape[0] * Loc3.data.shape[1] * Loc3.data.shape[2] * common_params.num_boxes[2], int(Loc3.data.shape[3] / common_params.num_boxes[2]))\n",
        "            Cls3 = Cls3.view(Cls3.data.shape[0] * Cls3.data.shape[1] * Cls3.data.shape[2] * common_params.num_boxes[2], int(Cls3.data.shape[3] / common_params.num_boxes[2]))\n",
        "\n",
        "            Loc4 = Loc4.view(Loc4.data.shape[0] * Loc4.data.shape[1] * Loc4.data.shape[2] * common_params.num_boxes[3], int(Loc4.data.shape[3] / common_params.num_boxes[3]))\n",
        "            Cls4 = Cls4.view(Cls4.data.shape[0] * Cls4.data.shape[1] * Cls4.data.shape[2] * common_params.num_boxes[3], int(Cls4.data.shape[3] / common_params.num_boxes[3]))\n",
        "\n",
        "            Loc5 = Loc5.view(Loc5.data.shape[0] * Loc5.data.shape[1] * Loc5.data.shape[2] * common_params.num_boxes[4], int(Loc5.data.shape[3] / common_params.num_boxes[4]))\n",
        "            Cls5 = Cls5.view(Cls5.data.shape[0] * Cls5.data.shape[1] * Cls5.data.shape[2] * common_params.num_boxes[4], int(Cls5.data.shape[3] / common_params.num_boxes[4]))\n",
        "\n",
        "            Loc6 = Loc6.view(Loc6.data.shape[0] * Loc6.data.shape[1] * Loc6.data.shape[2] * common_params.num_boxes[5], int(Loc6.data.shape[3] / common_params.num_boxes[5]))\n",
        "            Cls6 = Cls6.view(Cls6.data.shape[0] * Cls6.data.shape[1] * Cls6.data.shape[2] * common_params.num_boxes[5], int(Cls6.data.shape[3] / common_params.num_boxes[5]))\n",
        "\n",
        "            return (Loc1, Cls1, Loc2, Cls2, Loc3, Cls3, Loc4, Cls4, Loc5, Cls5, Loc6, Cls6)\n",
        "\n",
        "        else:\n",
        "            return (Loc1, Cls1, Loc2, Cls2, Loc3, Cls3, Loc4, Cls4, Loc5, Cls5, Loc6, Cls6)\n",
        "\n",
        "model = SSDNet()\n",
        "if use_cuda:\n",
        "  model.cuda()\n",
        "# モデルの情報を表示\n",
        "torchsummary.summary(model, (3, 300, 300))"
      ],
      "metadata": {
        "id": "I4sDBJJdzseq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S3qxW4OiL5Hs"
      },
      "source": [
        "# 損失関数の定義\n",
        "学習に使用する損失関数と最適化手法を定義します．Loss計算はGPUよりもCPUの方が高速に処理できるため，CPU上で処理を行います．\\\n",
        "最初に各層のクラス推定と位置推定を行う層と教師データのBounding box情報を照らし合わせるために，同等の特徴マップを用意して教示データを格納しています．\n",
        "クラス推定にはクロスエントロピー誤差，位置推定にはsmooth L1誤差を用いて誤差を算出します．"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class MYLoss(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(MYLoss, self).__init__()\n",
        "\n",
        "    def forward(self, Loc, Cls, gt_box_batch, df_box_batch, idx_batch, cls_batch, bat_s, mining):\n",
        "        if mining:\n",
        "            # hard negative mining有効時のクラスラベル\n",
        "            cls_t1 = torch.ones((bat_s, common_params.num_boxes[0], common_params.map_sizes[0], common_params.map_sizes[0]), dtype = torch.int64).cpu() * -1\n",
        "            cls_t2 = torch.ones((bat_s, common_params.num_boxes[1], common_params.map_sizes[1], common_params.map_sizes[1]), dtype = torch.int64).cpu() * -1\n",
        "            cls_t3 = torch.ones((bat_s, common_params.num_boxes[2], common_params.map_sizes[2], common_params.map_sizes[2]), dtype = torch.int64).cpu() * -1\n",
        "            cls_t4 = torch.ones((bat_s, common_params.num_boxes[3], common_params.map_sizes[3], common_params.map_sizes[3]), dtype = torch.int64).cpu() * -1\n",
        "            cls_t5 = torch.ones((bat_s, common_params.num_boxes[4], common_params.map_sizes[4], common_params.map_sizes[4]), dtype = torch.int64).cpu() * -1\n",
        "            cls_t6 = torch.ones((bat_s, common_params.num_boxes[5], common_params.map_sizes[5], common_params.map_sizes[5]), dtype = torch.int64).cpu() * -1\n",
        "        else:\n",
        "            # hard negative mining無効時のクラスラベル\n",
        "            cls_t1 = torch.zeros((bat_s, common_params.num_boxes[0], common_params.map_sizes[0], common_params.map_sizes[0]), dtype = torch.int64).cpu()\n",
        "            cls_t2 = torch.zeros((bat_s, common_params.num_boxes[1], common_params.map_sizes[1], common_params.map_sizes[1]), dtype = torch.int64).cpu()\n",
        "            cls_t3 = torch.zeros((bat_s, common_params.num_boxes[2], common_params.map_sizes[2], common_params.map_sizes[2]), dtype = torch.int64).cpu()\n",
        "            cls_t4 = torch.zeros((bat_s, common_params.num_boxes[3], common_params.map_sizes[3], common_params.map_sizes[3]), dtype = torch.int64).cpu()\n",
        "            cls_t5 = torch.zeros((bat_s, common_params.num_boxes[4], common_params.map_sizes[4], common_params.map_sizes[4]), dtype = torch.int64).cpu()\n",
        "            cls_t6 = torch.zeros((bat_s, common_params.num_boxes[5], common_params.map_sizes[5], common_params.map_sizes[5]), dtype = torch.int64).cpu()\n",
        "\n",
        "        # bounding boxのオフセットベクトルの教示データ\n",
        "        loc_t1 = torch.zeros((bat_s, common_params.num_boxes[0] * common_params.num_of_offset_dims, common_params.map_sizes[0], common_params.map_sizes[0]), dtype = torch.float32).cpu()\n",
        "        loc_t2 = torch.zeros((bat_s, common_params.num_boxes[1] * common_params.num_of_offset_dims, common_params.map_sizes[1], common_params.map_sizes[1]), dtype = torch.float32).cpu()\n",
        "        loc_t3 = torch.zeros((bat_s, common_params.num_boxes[2] * common_params.num_of_offset_dims, common_params.map_sizes[2], common_params.map_sizes[2]), dtype = torch.float32).cpu()\n",
        "        loc_t4 = torch.zeros((bat_s, common_params.num_boxes[3] * common_params.num_of_offset_dims, common_params.map_sizes[3], common_params.map_sizes[3]), dtype = torch.float32).cpu()\n",
        "        loc_t5 = torch.zeros((bat_s, common_params.num_boxes[4] * common_params.num_of_offset_dims, common_params.map_sizes[4], common_params.map_sizes[4]), dtype = torch.float32).cpu()\n",
        "        loc_t6 = torch.zeros((bat_s, common_params.num_boxes[5] * common_params.num_of_offset_dims, common_params.map_sizes[5], common_params.map_sizes[5]), dtype = torch.float32).cpu()\n",
        "\n",
        "        cls_t = ([cls_t1, cls_t2, cls_t3, cls_t4, cls_t5, cls_t5])\n",
        "        loc_t = ([loc_t1, loc_t2, loc_t3, loc_t4, loc_t5, loc_t6])\n",
        "\n",
        "        for b in range(0, len(idx_batch)):\n",
        "            for i in range(0, len(idx_batch[b])):\n",
        "                if int(idx_batch[b][i][1]) == -100: break\n",
        "                fmap_layer = idx_batch[b][i][1]\n",
        "                fmap_position = idx_batch[b][i][2]\n",
        "                df_box_num = idx_batch[b][i][3]\n",
        "                st_box_idx = df_box_num * common_params.num_of_offset_dims\n",
        "                ed_box_idx = st_box_idx + common_params.num_of_offset_dims\n",
        "\n",
        "                c = int(fmap_position % common_params.map_sizes[fmap_layer])\n",
        "                r = int(fmap_position / common_params.map_sizes[fmap_layer])\n",
        "\n",
        "                item_class_id = cls_batch[b][i]\n",
        "\n",
        "                # 1〜6番目のdefault boxのクラスとオフセットの教示データを格納\n",
        "                gt_box_batch_idx = gt_box_batch[b][i]\n",
        "                df_box_batch_idx = df_box_batch[b][i]\n",
        "\n",
        "                cls_t[fmap_layer][b, df_box_num, r, c] = cls_batch[b][i]\n",
        "                loc_t[fmap_layer][b, st_box_idx : ed_box_idx, r, c] = (gt_box_batch_idx - df_box_batch_idx) / common_params.loc_var\n",
        "\n",
        "        # 1〜6階層目の教示confidence mapの次元を(バッチ数, DF box数, 高さ, 幅)から(バッチ数, 高さ, 幅, DF box数)に転置\n",
        "        cls_t1 = cls_t1.permute(0, 2, 3, 1).contiguous()\n",
        "        cls_t2 = cls_t2.permute(0, 2, 3, 1).contiguous()\n",
        "        cls_t3 = cls_t3.permute(0, 2, 3, 1).contiguous()\n",
        "        cls_t4 = cls_t4.permute(0, 2, 3, 1).contiguous()\n",
        "        cls_t5 = cls_t5.permute(0, 2, 3, 1).contiguous()\n",
        "        cls_t6 = cls_t6.permute(0, 2, 3, 1).contiguous()\n",
        "\n",
        "        # 1〜6階層目の教示confidence mapの各次元数を(バッチ数, 高さ, 幅, DF box数)から(バッチ数 * 高さ * 幅 * DF box数)にreshape\n",
        "        cls_t1 = cls_t1.view(cls_t1.data.shape[0] * cls_t1.data.shape[1] * cls_t1.data.shape[2] * common_params.num_boxes[0])\n",
        "        cls_t2 = cls_t2.view(cls_t2.data.shape[0] * cls_t2.data.shape[1] * cls_t2.data.shape[2] * common_params.num_boxes[1])\n",
        "        cls_t3 = cls_t3.view(cls_t3.data.shape[0] * cls_t3.data.shape[1] * cls_t3.data.shape[2] * common_params.num_boxes[2])\n",
        "        cls_t4 = cls_t4.view(cls_t4.data.shape[0] * cls_t4.data.shape[1] * cls_t4.data.shape[2] * common_params.num_boxes[3])\n",
        "        cls_t5 = cls_t5.view(cls_t5.data.shape[0] * cls_t5.data.shape[1] * cls_t5.data.shape[2] * common_params.num_boxes[4])\n",
        "        cls_t6 = cls_t6.view(cls_t6.data.shape[0] * cls_t6.data.shape[1] * cls_t6.data.shape[2] * common_params.num_boxes[5])\n",
        "\n",
        "        # 1〜6階層目の教示localization mapの次元を(バッチ数, オフセット次元数 * DF box数, 高さ, 幅)から(バッチ数, 高さ, 幅, オフセット次元数 * DF box数)に転置\n",
        "        loc_t1 = loc_t1.permute(0, 2, 3, 1).contiguous()\n",
        "        loc_t2 = loc_t2.permute(0, 2, 3, 1).contiguous()\n",
        "        loc_t3 = loc_t3.permute(0, 2, 3, 1).contiguous()\n",
        "        loc_t4 = loc_t4.permute(0, 2, 3, 1).contiguous()\n",
        "        loc_t5 = loc_t5.permute(0, 2, 3, 1).contiguous()\n",
        "        loc_t6 = loc_t6.permute(0, 2, 3, 1).contiguous()\n",
        "\n",
        "        # 1〜6階層目の教示localization mapの各次元数を(バッチ数, 高さ, 幅, オフセット次元数 * DF box数)から(バッチ数 * 高さ * 幅 * DF box数, オフセット次元数)にreshape\n",
        "        loc_t1 = loc_t1.view(loc_t1.data.shape[0] * loc_t1.data.shape[1] * loc_t1.data.shape[2] * common_params.num_boxes[0], int(loc_t1.data.shape[3] / common_params.num_boxes[0]))\n",
        "        loc_t2 = loc_t2.view(loc_t2.data.shape[0] * loc_t2.data.shape[1] * loc_t2.data.shape[2] * common_params.num_boxes[1], int(loc_t2.data.shape[3] / common_params.num_boxes[1]))\n",
        "        loc_t3 = loc_t3.view(loc_t3.data.shape[0] * loc_t3.data.shape[1] * loc_t3.data.shape[2] * common_params.num_boxes[2], int(loc_t3.data.shape[3] / common_params.num_boxes[2]))\n",
        "        loc_t4 = loc_t4.view(loc_t4.data.shape[0] * loc_t4.data.shape[1] * loc_t4.data.shape[2] * common_params.num_boxes[3], int(loc_t4.data.shape[3] / common_params.num_boxes[3]))\n",
        "        loc_t5 = loc_t5.view(loc_t5.data.shape[0] * loc_t5.data.shape[1] * loc_t5.data.shape[2] * common_params.num_boxes[4], int(loc_t5.data.shape[3] / common_params.num_boxes[4]))\n",
        "        loc_t6 = loc_t6.view(loc_t6.data.shape[0] * loc_t6.data.shape[1] * loc_t6.data.shape[2] * common_params.num_boxes[5], int(loc_t6.data.shape[3] / common_params.num_boxes[5]))\n",
        "\n",
        "        # 1〜6階層目の教示confidence mapを結合\n",
        "        Cls_T = torch.cat([cls_t1, cls_t2, cls_t3, cls_t4, cls_t5, cls_t6], dim = 0)\n",
        "\n",
        "        # 1〜6階層目の教示localization mapを結合\n",
        "        Loc_T = torch.cat([loc_t1, loc_t2, loc_t3, loc_t4, loc_t5, loc_t6], dim = 0)\n",
        "\n",
        "        # loss計算\n",
        "        x_entropy = nn.CrossEntropyLoss(ignore_index=-1)\n",
        "        smooth_L1 = nn.SmoothL1Loss(reduction='none')\n",
        "\n",
        "        if use_cuda:\n",
        "          Cls = Cls.cuda()\n",
        "          Cls_T = Cls_T.cuda()\n",
        "          Loc = Loc.cuda()\n",
        "          Loc_T = Loc_T.cuda()\n",
        "\n",
        "        # confidence mapのloss\n",
        "        loss_cls = x_entropy(Cls, Cls_T)\n",
        "        # localization mapのloss\n",
        "        loss_loc = smooth_L1(Loc, Loc_T)\n",
        "        loss_loc = torch.sum(loss_loc, dim=-1)\n",
        "        positive_samples = torch.tensor(Cls_T > 0, dtype=torch.float32)\n",
        "        if use_cuda:\n",
        "          positive_samples = positive_samples.cuda()\n",
        "        loss_loc *= positive_samples\n",
        "        n_positives = positive_samples.sum()\n",
        "        loss_loc = torch.sum(loss_loc) / n_positives\n",
        "\n",
        "        return loss_cls, loss_loc"
      ],
      "metadata": {
        "id": "i-AIraCSYV2f"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ELuHJ5lGOmmv"
      },
      "source": [
        "# 学習\n",
        "損失関数の定義まで操作が完了したら，学習する段階に入ります．"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XmE767JNUe9s"
      },
      "source": [
        "## ネットワークと最適化方法の設定\n",
        "先ほど作成したネットワーク(SSDNet)を`.cuda()`によりGPU上に展開します．GPUにネットワークモデルをのせることで，並列演算の得意なGPUによる高速な学習が可能です．\\\n",
        "学習を行う際の最適化方法としてSGD（確率的勾配降下法）を利用します．\n",
        "定義済みのLoss計算(MYLoss)を損失関数として設定します．\n",
        "また，学習率を0.001，モーメンタムを0.9，重み減衰率0.0005をとして引数に与えます．"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "LEARNING_RATE = 10e-3\n",
        "MOMENTUM = 0.9\n",
        "WEIGHT_DECAY = 0.0005\n",
        "\n",
        "# SSDNetの読み込み\n",
        "ssd_model = SSDNet()\n",
        "if use_cuda:\n",
        "  ssd_model.cuda()\n",
        "\n",
        "# Setup optimizer\n",
        "optimizer = optim.SGD(ssd_model.parameters(), lr = LEARNING_RATE, momentum = MOMENTUM, weight_decay = WEIGHT_DECAY)\n",
        "\n",
        "# loss function\n",
        "loss_function = MYLoss()"
      ],
      "metadata": {
        "id": "Wuuk3ztZYYEd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ds56MjEnUqFA"
      },
      "source": [
        "## モデルの学習\n",
        "ネットワークと最適化方法の設定完了後，学習を開始します．\n",
        "学習時は，入力された画像の推論が実行され，confidense mapとlocalization mapのlossの集計して，誤差逆伝播によるパラメータの更新が行われます．\\\n",
        "物体検出の学習には時間を要するため，エポック数は30とします．学習には20分ほどかかります．\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# エポック数の設定\n",
        "epoch_num = 30\n",
        "\n",
        "all_iter = 0\n",
        "num_report_iter = 50\n",
        "\n",
        "begin_at = time.time()\n",
        "for epoch in range(1, epoch_num+1):\n",
        "    running_loss_cls = 0.0\n",
        "    running_loss_loc = 0.0\n",
        "    running_loss = 0.0\n",
        "\n",
        "    # epochが4回に1回の割合でhard negative mining有効\n",
        "    mining = True if (epoch) % 4 == 0 else False\n",
        "\n",
        "    now_datetime = datetime.datetime.today()\n",
        "    now_datetime = str(now_datetime.year) + '-' + str('%02d' % now_datetime.month) + '-' + str('%02d' % now_datetime.day) + '@' + str('%02d' % now_datetime.hour) + ':' + str('%02d' % now_datetime.minute) + ':' + str('%02d' % now_datetime.second)\n",
        "\n",
        "    for num_batch, data in enumerate(train_dataloader):\n",
        "        all_iter += 1\n",
        "\n",
        "        img_batch, gt_box_batch, df_box_batch, idx_batch, cls_batch, conf_img_batch = data\n",
        "\n",
        "        train_img = img_batch\n",
        "        if use_cuda:\n",
        "            train_img = train_img.cuda()\n",
        "\n",
        "        # SSD net forward\n",
        "        Loc1, Cls1, Loc2, Cls2, Loc3, Cls3, Loc4, Cls4, Loc5, Cls5, Loc6, Cls6 = ssd_model(train_img, train=True)\n",
        "\n",
        "        # ネットワークから出力されたconfidence mapを1〜6階層目まで結合\n",
        "        Loc = torch.cat([Loc1, Loc2, Loc3, Loc4, Loc5, Loc6], dim = 0)\n",
        "        # ネットワークから出力されたlocalization mapを1〜6階層目まで結合\n",
        "        Cls = torch.cat([Cls1, Cls2, Cls3, Cls4, Cls5, Cls6], dim = 0)\n",
        "\n",
        "        # lossを計算\n",
        "        loss_cls, loss_loc = loss_function(Loc, Cls, gt_box_batch, df_box_batch, idx_batch, cls_batch, batch_size, mining)\n",
        "        loss = (loss_cls * 0.5 + loss_loc * 0.5)\n",
        "\n",
        "        optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        elapsed_time = datetime.timedelta(seconds = time.time() - begin_at)\n",
        "        sys.stdout.write(\"\\rEpoch: {}, Iter: {}, Loss:{} (cls: {}, loc: {})\".format(epoch, all_iter, loss.data, loss_cls.data, loss_loc.data))\n",
        "        sys.stdout.flush()\n",
        "        time.sleep(0.01)\n",
        "\n",
        "        running_loss_cls += loss_cls\n",
        "        running_loss_loc += loss_loc\n",
        "        running_loss += loss\n",
        "\n",
        "    running_loss /= all_iter\n",
        "    running_loss_cls /= all_iter\n",
        "    running_loss_loc /= all_iter\n",
        "    print(\"\\r\\r\\r[Epoch:{} Iter:{}] Time: {}, Ave_Loss:{} (cls: {}, loc: {}) mining: {}\".format(epoch, all_iter, elapsed_time, running_loss, running_loss_cls, running_loss_loc, mining))\n",
        "\n",
        "print(\"\\nExit Training\\n\")"
      ],
      "metadata": {
        "id": "siHwMvVnzm5J"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KEH7qzvi5RJE"
      },
      "source": [
        "# 評価\n",
        "先程学習したモデルを使って評価を行ってみましょう．\\\n",
        "MODEL_PATHに学習済みモデルを指定して実行します．\n",
        "argsクラス内で定義されている`IN_DIR`は評価対象のフォルダ，`MODEL_PATH`は学習済みモデルを指定します．\n",
        "以下の評価コードを実行することで，物体検出後の画像とそのバウンディングボックス情報が保存されます．\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "IN_DIR = \"./val/rgb\"                #  評価データのパス\n",
        "OUT_DIR = \"./out\"                   #  結果保存先\n",
        "SAVE_DATA = True                    #  画像を保存\n",
        "\n",
        "# クラスラベル (クラス名にはスペース(空白)は禁止)\n",
        "labels = common_params.arc_labels\n",
        "class_color = common_params.arc_class_color[:, ::-1]\n",
        "print(len(labels))\n",
        "\n",
        "# confidence mapのsoftmaxを計算\n",
        "def mboxSoftmax(confidence_maps, num_classes, num_boxes):\n",
        "    s = np.zeros((confidence_maps.shape[0], confidence_maps.shape[1], confidence_maps.shape[2]), np.float32)\n",
        "\n",
        "    bs = 0\n",
        "    be = num_classes\n",
        "    for b in range(0, num_boxes):\n",
        "        t = confidence_maps.detach().numpy()[bs : be, :, :]\n",
        "\n",
        "        total = 0\n",
        "        for i in range(0, t.shape[0]):\n",
        "            total += np.exp(t[i, :, :])\n",
        "\n",
        "        for i in range(0, t.shape[0]):\n",
        "            s[bs + i, :, :] = np.exp(t[i, :, :]) / total\n",
        "\n",
        "        bs = be\n",
        "        be += num_classes\n",
        "    return s\n",
        "\n",
        "# クラス確率の高いdefault boxを検出\n",
        "def multiBoxDetection(cls_score_maps, localization_maps, num_dbox, num_class, offset_dim, min_size, max_size, step, aspect_ratio):\n",
        "    box_offsets = []\n",
        "    default_boxes = []\n",
        "    class_labels = []\n",
        "    class_scores = []\n",
        "\n",
        "    img_width = common_params.insize\n",
        "    img_height = common_params.insize\n",
        "\n",
        "    map_size = cls_score_maps.shape[1] * cls_score_maps.shape[2]\n",
        "    for i in range(0, map_size):\n",
        "\n",
        "        c = int(i % cls_score_maps.shape[1])\n",
        "        r = int(i / cls_score_maps.shape[1])\n",
        "\n",
        "        mbox_max_val = 0\n",
        "        mbox_max_idx = 0\n",
        "        mbox_num = 0\n",
        "\n",
        "        bs = 0\n",
        "        be = num_class\n",
        "        for b in range(0, num_dbox):\n",
        "\n",
        "            max_val = np.max(cls_score_maps[bs : be, r, c])\n",
        "            max_idx = int(np.argmax(cls_score_maps[bs : be, r, c]))\n",
        "\n",
        "            if max_val > mbox_max_val and max_idx != 0:\n",
        "                mbox_max_val = max_val\n",
        "                mbox_max_idx = max_idx\n",
        "                mbox_num = b\n",
        "\n",
        "            bs = be\n",
        "            be += num_class\n",
        "\n",
        "        bs = mbox_num * offset_dim\n",
        "        be = bs + offset_dim\n",
        "        b_offset = localization_maps[bs : be, r, c]\n",
        "\n",
        "        offset_ = 0.5\n",
        "        if mbox_max_val >= 0.1: # pyt0.7  cha0.5\n",
        "            center_x = float((c + offset_) * step)\n",
        "            center_y = float((r + offset_) * step)\n",
        "\n",
        "            if mbox_num == 0:\n",
        "                box_width = box_height = min_size\n",
        "                xmin = (center_x - box_width / 2.) / img_width\n",
        "                ymin = (center_y - box_height / 2.) / img_height\n",
        "                xmax = (center_x + box_width / 2.) / img_width\n",
        "                ymax = (center_y + box_height / 2.) / img_height\n",
        "            elif mbox_num == 1:\n",
        "                box_width = box_height = np.sqrt(min_size * max_size)\n",
        "                xmin = (center_x - box_width / 2.) / img_width\n",
        "                ymin = (center_y - box_height / 2.) / img_height\n",
        "                xmax = (center_x + box_width / 2.) / img_width\n",
        "                ymax = (center_y + box_height / 2.) / img_height\n",
        "            elif mbox_num == 2:\n",
        "                box_width = min_size * np.sqrt(float(aspect_ratio[0]))\n",
        "                box_height = min_size / np.sqrt(float(aspect_ratio[0]))\n",
        "                xmin = (center_x - box_width / 2.) / img_width\n",
        "                ymin = (center_y - box_height / 2.) / img_height\n",
        "                xmax = (center_x + box_width / 2.) / img_width\n",
        "                ymax = (center_y + box_height / 2.) / img_height\n",
        "            elif mbox_num == 3:\n",
        "                box_width = min_size * np.sqrt(1. / float(aspect_ratio[0]))\n",
        "                box_height = min_size / np.sqrt(1. / float(aspect_ratio[0]))\n",
        "                xmin = (center_x - box_width / 2.) / img_width\n",
        "                ymin = (center_y - box_height / 2.) / img_height\n",
        "                xmax = (center_x + box_width / 2.) / img_width\n",
        "                ymax = (center_y + box_height / 2.) / img_height\n",
        "            elif mbox_num == 4:\n",
        "                box_width = min_size * np.sqrt(float(aspect_ratio[1]))\n",
        "                box_height = min_size / np.sqrt(float(aspect_ratio[1]))\n",
        "                xmin = (center_x - box_width / 2.) / img_width\n",
        "                ymin = (center_y - box_height / 2.) / img_height\n",
        "                xmax = (center_x + box_width / 2.) / img_width\n",
        "                ymax = (center_y + box_height / 2.) / img_height\n",
        "            elif mbox_num == 5:\n",
        "                box_width = min_size * np.sqrt(1. / float(aspect_ratio[1]))\n",
        "                box_height = min_size / np.sqrt(1. / float(aspect_ratio[1]))\n",
        "                xmin = (center_x - box_width / 2.) / img_width\n",
        "                ymin = (center_y - box_height / 2.) / img_height\n",
        "                xmax = (center_x + box_width / 2.) / img_width\n",
        "                ymax = (center_y + box_height / 2.) / img_height\n",
        "\n",
        "            box_offsets.append(b_offset)\n",
        "            default_boxes.append([min(max(xmin, 0.), 1.), min(max(ymin, 0.), 1.), min(max(xmax, 0.), 1.), min(max(ymax, 0.), 1.), mbox_num])\n",
        "            class_labels.append(mbox_max_idx)\n",
        "            class_scores.append(mbox_max_val)\n",
        "\n",
        "    return (box_offsets, default_boxes, class_labels, class_scores)\n",
        "\n",
        "# bounding box候補を検出 (default boxの補正)\n",
        "def candidatesDetection(offsets, default_boxes, class_labels, class_scores, num_classes, color_img, variance):\n",
        "\n",
        "    img_width = color_img.shape[1]\n",
        "    img_height = color_img.shape[0]\n",
        "\n",
        "    candidates = []\n",
        "    for i in range(0, num_classes):\n",
        "        candidates.append([])\n",
        "\n",
        "    for det in range(0, len(class_labels)):\n",
        "\n",
        "        pred_xmin = (default_boxes[det][0] + offsets[det][0] * variance)\n",
        "        pred_ymin = (default_boxes[det][1] + offsets[det][1] * variance)\n",
        "        pred_xmax = (default_boxes[det][2] + offsets[det][2] * variance)\n",
        "        pred_ymax = (default_boxes[det][3] + offsets[det][3] * variance)\n",
        "\n",
        "        pred_xmin = min(max(pred_xmin, 0.), 1.) * img_width\n",
        "        pred_ymin = min(max(pred_ymin, 0.), 1.) * img_height\n",
        "        pred_xmax = min(max(pred_xmax, 0.), 1.) * img_width\n",
        "        pred_ymax = min(max(pred_ymax, 0.), 1.) * img_height\n",
        "        candidates[class_labels[det]].append([pred_xmin, pred_ymin, pred_xmax, pred_ymax, class_scores[det]])\n",
        "\n",
        "    return candidates\n",
        "\n",
        "# 同じクラスに属するbounding boxの重なり率を計算\n",
        "def jaccardOverlap(bbox1, bbox2):\n",
        "\n",
        "    if (bbox2[0] > bbox1[2]) or (bbox2[2] < bbox1[0]) or (bbox2[1] > bbox1[3]) or (bbox2[3] < bbox1[1]):\n",
        "        overlap = 0.\n",
        "    else:\n",
        "        inter_xmin = max(bbox1[0], bbox2[0])\n",
        "        inter_ymin = max(bbox1[1], bbox2[1])\n",
        "        inter_xmax = min(bbox1[2], bbox2[2])\n",
        "        inter_ymax = min(bbox1[3], bbox2[3])\n",
        "\n",
        "        inter_width = inter_xmax - inter_xmin\n",
        "        inter_height = inter_ymax - inter_ymin\n",
        "        inter_size = inter_width * inter_height\n",
        "\n",
        "        bbox1_size = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])\n",
        "        bbox2_size = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1])\n",
        "\n",
        "        overlap = inter_size / (bbox1_size + bbox2_size - inter_size)\n",
        "\n",
        "    return overlap\n",
        "\n",
        "# bounding box候補のnon-maximum suppresion\n",
        "def nonMaximumSuppresion(candidates):\n",
        "\n",
        "    overlap_th = 0.2\n",
        "\n",
        "    for i in range(0, len(candidates)):\n",
        "\n",
        "        box_num = len(candidates[i])\n",
        "\n",
        "        js = 0\n",
        "        for j in range(0, box_num):\n",
        "            ks = js + 1\n",
        "            for k in range(j + 1, box_num):\n",
        "\n",
        "                if ks >= len(candidates[i]) or js >= len(candidates[i]):\n",
        "                    continue\n",
        "\n",
        "                overlap = jaccardOverlap(candidates[i][js], candidates[i][ks])\n",
        "\n",
        "                if overlap >= overlap_th and candidates[i][js][4] >= candidates[i][ks][4]:\n",
        "                    candidates[i].pop(ks)\n",
        "                    ks -= 1\n",
        "                elif overlap >= overlap_th and candidates[i][js][4] < candidates[i][ks][4]:\n",
        "                    candidates[i][js], candidates[i][ks] = candidates[i][ks], candidates[i][js]\n",
        "                    candidates[i].pop(ks)\n",
        "                    ks -= 1\n",
        "\n",
        "                ks += 1\n",
        "            js += 1\n",
        "\n",
        "    return candidates\n",
        "\n",
        "# box検出結果を画像に描画してテキストとともに保存\n",
        "def saveDetection(final_detections, out_img, filename):\n",
        "    # 検出したbounding boxを画像に描画\n",
        "    offset_ = 0.5\n",
        "    font = cv.FONT_HERSHEY_SIMPLEX\n",
        "    for i in range(0, len(final_detections)):\n",
        "            class_name = labels[i]\n",
        "            color = class_color[i]\n",
        "            for j in range(0, len(final_detections[i])):\n",
        "                p1 = int(final_detections[i][j][0] + offset_)\n",
        "                p2 = int(final_detections[i][j][1] + offset_)\n",
        "                p3 = int(final_detections[i][j][2] + offset_)\n",
        "                p4 = int(final_detections[i][j][3] + offset_)\n",
        "                colors = (int(color[0]), int(color[1]), int(color[2]))\n",
        "                cv.rectangle(out_img, (p1, p2), (p3, p4), colors, 5)\n",
        "                q1 = p1\n",
        "                q2 = p4\n",
        "                # クラス名表示\n",
        "                cv.rectangle(out_img, (q1, q2 - 30), (q1 + len(class_name) * 25, q2), colors, -1)\n",
        "                cv.putText(out_img, str(i) + \" \" + class_name + \": \" + str(('%.2f' % final_detections[i][j][4])),\n",
        "                           (q1, q2 - 8), font, 0.7, (0, 0, 0), 2, cv.LINE_AA)\n",
        "                cv.putText(out_img, str(i) + \" \" + class_name + \": \" + str(('%.2f' % final_detections[i][j][4])),\n",
        "                           (q1, q2 - 8), font, 0.7, (255, 255, 255), 1, cv.LINE_AA)\n",
        "                # 検出したクラスのラベル、スコア、座標を出力\n",
        "                if SAVE_DATA:\n",
        "                    f = open(OUT_DIR + '/detection_txt/' + filename + 'res.txt', 'a')\n",
        "                    f.write(\"{} {} {} {} {} {}\\n\".format(str(i), final_detections[i][j][4], p1, p2, p3, p4))\n",
        "                    f.close()\n",
        "    # 画像保存\n",
        "    if SAVE_DATA:\n",
        "        cv.imwrite(OUT_DIR + '/detection/' + filename + '.png', out_img)\n",
        "\n",
        "    return out_img\n",
        "\n",
        "def detection(img, ssd_model, filename, min_sizes, max_sizes):\n",
        "    # タイマーリセット\n",
        "    total_time = 0.0\n",
        "    processing_time = 0.0\n",
        "    drawing_time = 0.0\n",
        "    classes = []\n",
        "\n",
        "    # 入力画像をSSDの入力サイズにリサイズ\n",
        "    start = time.time()\n",
        "    input_img = cv.resize(img, (common_params.insize, common_params.insize), interpolation = cv.INTER_CUBIC)\n",
        "    out_img = img.copy()\n",
        "\n",
        "    input_img = input_img.astype(np.float32)\n",
        "    input_img -= np.array([103.939, 116.779, 123.68])\n",
        "    input_img = input_img.transpose(2, 0, 1)\n",
        "    input_data = []\n",
        "    input_data.append(input_img)\n",
        "\n",
        "    x_data = torch.tensor(input_data, dtype=torch.float32)\n",
        "    if use_cuda:\n",
        "      x_data = x_data.cuda()\n",
        "    elapsed_time = time.time() - start\n",
        "    processing_time += elapsed_time\n",
        "    total_time += elapsed_time\n",
        "\n",
        "    # SSDのforward\n",
        "    start = time.time()\n",
        "    Loc1, Cls1, Loc2, Cls2, Loc3, Cls3, Loc4, Cls4, Loc5, Cls5, Loc6, Cls6 = ssd_model(x_data, train=False)\n",
        "    elapsed_time = time.time() - start\n",
        "    processing_time += elapsed_time\n",
        "    total_time += elapsed_time\n",
        "\n",
        "    # CPUで処理\n",
        "    start = time.time()\n",
        "    Loc1, Cls1 = Loc1.cpu(), Cls1.cpu()\n",
        "    Loc2, Cls2 = Loc2.cpu(), Cls2.cpu()\n",
        "    Loc3, Cls3 = Loc3.cpu(), Cls3.cpu()\n",
        "    Loc4, Cls4 = Loc4.cpu(), Cls4.cpu()\n",
        "    Loc5, Cls5 = Loc5.cpu(), Cls5.cpu()\n",
        "    Loc6, Cls6 = Loc6.cpu(), Cls6.cpu()\n",
        "    elapsed_time = time.time() - start\n",
        "    processing_time += elapsed_time\n",
        "    total_time += elapsed_time\n",
        "\n",
        "    # 各階層のconfidence mapのsoftmaxを計算 (clsの中身が怪しいそれかsoftmax)\n",
        "    start = time.time()\n",
        "    cls_score1 = mboxSoftmax(Cls1[0], common_params.num_of_classes, common_params.num_boxes[0])\n",
        "    cls_score2 = mboxSoftmax(Cls2[0], common_params.num_of_classes, common_params.num_boxes[1])\n",
        "    cls_score3 = mboxSoftmax(Cls3[0], common_params.num_of_classes, common_params.num_boxes[2])\n",
        "    cls_score4 = mboxSoftmax(Cls4[0], common_params.num_of_classes, common_params.num_boxes[3])\n",
        "    cls_score5 = mboxSoftmax(Cls5[0], common_params.num_of_classes, common_params.num_boxes[4])\n",
        "    cls_score6 = mboxSoftmax(Cls6[0], common_params.num_of_classes, common_params.num_boxes[5])\n",
        "    elapsed_time = time.time() - start\n",
        "    processing_time += elapsed_time\n",
        "    total_time += elapsed_time\n",
        "\n",
        "    # クラス確率の高いdefault boxの検出\n",
        "    Loc1 = Loc1.detach().numpy()\n",
        "    Loc2 = Loc2.detach().numpy()\n",
        "    Loc3 = Loc3.detach().numpy()\n",
        "    Loc4 = Loc4.detach().numpy()\n",
        "    Loc5 = Loc5.detach().numpy()\n",
        "    Loc6 = Loc6.detach().numpy()\n",
        "\n",
        "    start = time.time()\n",
        "    offsets1, default_boxes1, class_labels1, class_scores1 = multiBoxDetection(cls_score1, Loc1[0], common_params.num_boxes[0], common_params.num_of_classes, common_params.num_of_offset_dims, min_sizes[0], max_sizes[0], common_params.steps[0], common_params.aspect_ratios[0])\n",
        "    offsets2, default_boxes2, class_labels2, class_scores2 = multiBoxDetection(cls_score2, Loc2[0], common_params.num_boxes[1], common_params.num_of_classes, common_params.num_of_offset_dims, min_sizes[1], max_sizes[1], common_params.steps[1], common_params.aspect_ratios[1])\n",
        "    offsets3, default_boxes3, class_labels3, class_scores3 = multiBoxDetection(cls_score3, Loc3[0], common_params.num_boxes[2], common_params.num_of_classes, common_params.num_of_offset_dims, min_sizes[2], max_sizes[2], common_params.steps[2], common_params.aspect_ratios[2])\n",
        "    offsets4, default_boxes4, class_labels4, class_scores4 = multiBoxDetection(cls_score4, Loc4[0], common_params.num_boxes[3], common_params.num_of_classes, common_params.num_of_offset_dims, min_sizes[3], max_sizes[3], common_params.steps[3], common_params.aspect_ratios[3])\n",
        "    offsets5, default_boxes5, class_labels5, class_scores5 = multiBoxDetection(cls_score5, Loc5[0], common_params.num_boxes[4], common_params.num_of_classes, common_params.num_of_offset_dims, min_sizes[4], max_sizes[4], common_params.steps[4], common_params.aspect_ratios[4])\n",
        "    offsets6, default_boxes6, class_labels6, class_scores6 = multiBoxDetection(cls_score6, Loc6[0], common_params.num_boxes[5], common_params.num_of_classes, common_params.num_of_offset_dims, min_sizes[5], max_sizes[5], common_params.steps[5], common_params.aspect_ratios[5])\n",
        "    elapsed_time = time.time() - start\n",
        "    processing_time += elapsed_time\n",
        "    total_time += elapsed_time\n",
        "\n",
        "    # オフセットベクトルによりdefault boxを補正\n",
        "    start = time.time()\n",
        "    candidates1 = candidatesDetection(offsets1, default_boxes1, class_labels1, class_scores1, common_params.num_of_classes, img, common_params.loc_var)\n",
        "    candidates2 = candidatesDetection(offsets2, default_boxes2, class_labels2, class_scores2, common_params.num_of_classes, img, common_params.loc_var)\n",
        "    candidates3 = candidatesDetection(offsets3, default_boxes3, class_labels3, class_scores3, common_params.num_of_classes, img, common_params.loc_var)\n",
        "    candidates4 = candidatesDetection(offsets4, default_boxes4, class_labels4, class_scores4, common_params.num_of_classes, img, common_params.loc_var)\n",
        "    candidates5 = candidatesDetection(offsets5, default_boxes5, class_labels5, class_scores5, common_params.num_of_classes, img, common_params.loc_var)\n",
        "    candidates6 = candidatesDetection(offsets6, default_boxes6, class_labels6, class_scores6, common_params.num_of_classes, img, common_params.loc_var)\n",
        "    elapsed_time = time.time() - start\n",
        "    processing_time += elapsed_time\n",
        "    total_time += elapsed_time\n",
        "\n",
        "    # 各階層のbounding box候補を統合\n",
        "    start = time.time()\n",
        "    all_candidate = []\n",
        "    for i in range(0, common_params.num_of_classes):\n",
        "        all_candidate.append([])\n",
        "        all_candidate[i].extend(candidates1[i])\n",
        "        all_candidate[i].extend(candidates2[i])\n",
        "        all_candidate[i].extend(candidates3[i])\n",
        "        all_candidate[i].extend(candidates4[i])\n",
        "        all_candidate[i].extend(candidates5[i])\n",
        "        all_candidate[i].extend(candidates6[i])\n",
        "    elapsed_time = time.time() - start\n",
        "    processing_time += elapsed_time\n",
        "    total_time += elapsed_time\n",
        "\n",
        "    # non-maximum suppresionによりbounding boxの最終結果を出力\n",
        "    start = time.time()\n",
        "    final_detections = nonMaximumSuppresion(all_candidate)\n",
        "    elapsed_time = time.time() - start\n",
        "    processing_time += elapsed_time\n",
        "    total_time += elapsed_time\n",
        "\n",
        "    # 画像保存\n",
        "    start = time.time()\n",
        "    out_img = saveDetection(final_detections, out_img, filename)\n",
        "    elapsed_time = time.time() - start\n",
        "    drawing_time += elapsed_time\n",
        "    total_time += elapsed_time\n",
        "\n",
        "    fps_end_time = time.time()\n",
        "    fps = 1 / total_time\n",
        "    print(\"Total Time : \", total_time)\n",
        "    # print(\"Processing Time : \", processing_time)\n",
        "    # print(\"FPS: \", fps)\n",
        "    print(\"----- Detection Done -----\")\n",
        "\n",
        "def test_out():\n",
        "    # ディレクトリの作成\n",
        "    if SAVE_DATA:\n",
        "        make_dir_list = [\n",
        "            OUT_DIR,\n",
        "            OUT_DIR + '/detection',                     # 検出結果のb-box画像\n",
        "            OUT_DIR + '/detection_txt'                  # 検出結果のb-boxテキスト\n",
        "        ]\n",
        "        for dirname in make_dir_list:\n",
        "            if not(os.path.exists(dirname)): os.mkdir(dirname)\n",
        "    # default boxのサイズリスト計算\n",
        "    step = int(math.floor((common_params.max_ratio - common_params.min_ratio) / (len(common_params.mbox_source_layers) - 2)))\n",
        "    min_sizes = []\n",
        "    max_sizes = []\n",
        "    for ratio in range(common_params.min_ratio, common_params.max_ratio + 1, step):\n",
        "            min_sizes.append(common_params.insize * ratio / 100.)\n",
        "            max_sizes.append(common_params.insize * (ratio + step) / 100.)\n",
        "    min_sizes = [common_params.insize * 10 / 100.] + min_sizes\n",
        "    max_sizes = [common_params.insize * 20 / 100.] + max_sizes\n",
        "\n",
        "    # 入力画像の読み込み\n",
        "    color_img = glob(os.path.join(IN_DIR, '*' + '.png'))\n",
        "    color_img.sort()\n",
        "    print(\"Evaluate images : \", len(color_img))\n",
        "\n",
        "    # 読み込んだリストを順次検出\n",
        "    ssd_model.eval()\n",
        "    for lf in range(len(color_img)):\n",
        "        print(\"{} / {}\".format(lf+1, len(color_img)))\n",
        "        img = cv.imread(color_img[lf])\n",
        "        filename, ext = os.path.splitext(os.path.basename(color_img[lf]))\n",
        "        detection(img, ssd_model, filename, min_sizes, max_sizes)\n",
        "\n",
        "if __name__ == '__main__':\n",
        "    test_out()"
      ],
      "metadata": {
        "id": "ACQxnuDszfxA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ffvI8wPyRJbQ"
      },
      "source": [
        "出力された検出結果の画像を確認してみましょう．"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!ls # ディレクトリの確認"
      ],
      "metadata": {
        "id": "TA9Z9Cq8zaC5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!ls out/"
      ],
      "metadata": {
        "id": "DyysPcMZzYuo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from IPython.display import Image,display_png\n",
        "display_png(Image('./out/detection/2017-016-2.png'))"
      ],
      "metadata": {
        "id": "yD5CqR9-zWlZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gd1eUsTxTN1S"
      },
      "source": [
        "BoundingBoxが表示されていなかったり，クラスが異なっていたり，うまく検出/識別できていないことが確認できます．\n",
        "この結果は，学習が30epochのみであり，十分に学習が行えていないためです．\n",
        "\n",
        "次に評価を実施します．検出されたバウンディングボックス1つに対して，その画像に存在するすべての教師信号のバウンディングボックスとIoUを比較して，最もIoUが高いものを探索します．見つかったボックスのクラスと検出されたボックスのクラスと2つのボックスのIoUを計算し，テキストファイルに書き出します．\n",
        "次に，書き出したテキストファイルを読んでConfusion Matrixを計算します．IoUが閾値以上かつクラスが合っていれば検出成功，そうでなければ検出失敗とします．計算が終了すると，画像形式のConfusion Matrixと，次の3つの値を出力します．\n",
        "- 識別率: 検出された物体のうち，クラスが正しかった割合\n",
        "- 未検出率: 全ての物体のうち，検出できなかった割合\n",
        "- 平均IoU: 全てのバウンディングボックスのIoUの平均"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%pwd\n",
        "!ls ./val/rgb/\n",
        "!ls ./out/detection_txt/"
      ],
      "metadata": {
        "id": "yA7Fv6qPzUJy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "image_path = \"./val/rgb/\"\n",
        "teach_path = \"./val/boundingbox/\"\n",
        "result_path = \"./out/detection_txt/\"\n",
        "\n",
        "# IoU matching results path\n",
        "match_result_path = result_path + \"matchingResults/\"\n",
        "\n",
        "# evaluation output(totalresult, confusion_matrix) path\n",
        "eval_result_path = result_path + \"eval/\"\n",
        "\n",
        "# wait time of cv2.waitKey (if 0 then no wait)\n",
        "WAITTIME = 0\n",
        "\n",
        "# if your detection results have a classlabel(e.g.: DVD, avery_binder...), set 1\n",
        "LABEL_FLAG = 1\n",
        "\n",
        "# if your detection results are normalized, set 1\n",
        "NORMALIZED = 0\n",
        "\n",
        "#IOU Threshold\n",
        "# IOU_THRESH = 0.55\n",
        "IOU_THRESH = 0.2\n",
        "\n",
        "# if teach labels have category and color classification results, set 1\n",
        "# normally need not change\n",
        "CAT_PASS_FLAG = 0\n",
        "\n",
        "# normally need not change\n",
        "THRESH = 0.35\n",
        "NCLASS = 40\n",
        "\n",
        "COLOR_TABLE = common_params.arc_class_color\n",
        "itemIDList = common_params.itemIDList\n",
        "\n",
        "def convNormalizedCord(data, height, width):\n",
        "    x = float(data[1]) * width\n",
        "    y = float(data[2]) * height\n",
        "    w = float(data[3]) * width\n",
        "    h = float(data[4]) * height\n",
        "    x1 = x - (w/2.)\n",
        "    y1 = y - (h/2.)\n",
        "    x2 = x + (w/2.)\n",
        "    y2 = y + (h/2.)\n",
        "    return(int(data[0]), int(x1), int(y1), int(x2), int(y2))\n",
        "\n",
        "def convResCord(data, height, width, normalized):\n",
        "    if normalized == 1:\n",
        "        data = convNormalizedCord(data, height, width)\n",
        "    if int(data[0]) == 0:\n",
        "        classID = 41\n",
        "    else:\n",
        "        classID = int(data[0])\n",
        "    return(classID, int(data[1]), int(data[2]), int(data[3]), int(data[4]))\n",
        "\n",
        "def getIOU(boxA, boxB):\n",
        "    # if the length of between boxAcenter and boxBcenter is too far, return 0\n",
        "    center_boxA = np.array([(boxA[0] + boxA[2]) / 2.0, (boxA[1] + boxA[3]) / 2.0])\n",
        "    center_boxB = np.array([(boxB[0] + boxB[2]) / 2.0, (boxB[1] + boxB[3]) / 2.0])\n",
        "    if np.linalg.norm(center_boxA - center_boxB) >= 500:\n",
        "        return 0\n",
        "\n",
        "    xA = max(boxA[0], boxB[0])\n",
        "    yA = max(boxA[1], boxB[1])\n",
        "    xB = min(boxA[2], boxB[2])\n",
        "    yB = min(boxA[3], boxB[3])\n",
        "\n",
        "    inter_area = (xB - xA + 1) * (yB - yA + 1)\n",
        "    boxA_area = (boxA[2] - boxA[0] + 1) * (boxA[3] - boxA[1] + 1)\n",
        "    boxB_area = (boxB[2] - boxB[0] + 1) * (boxB[3] - boxB[1] + 1)\n",
        "\n",
        "    iou = inter_area / float(boxA_area + boxB_area - inter_area)\n",
        "\n",
        "    return iou\n",
        "\n",
        "def readTxt(file_path, d_type, label_flag=0, cat_pass_flag=0):\n",
        "    coordinate = []\n",
        "    f = open(file_path, 'r')\n",
        "\n",
        "    if f != None:\n",
        "        for row in f:\n",
        "            data = row.split()\n",
        "            if(label_flag==1 and d_type == \"result\"):\n",
        "                data = [data[0], data[2], data[3], data[4], data[5]]\n",
        "            elif(d_type == \"result\"):\n",
        "                data = [data[0], data[1], data[2], data[3], data[4]]\n",
        "            elif(cat_pass_flag==1 and d_type == \"teach\"):\n",
        "                data = [data[0], data[3], data[4], data[5], data[6]]\n",
        "            elif(d_type == \"teach\" or d_type == \"evaluate\"):\n",
        "                data = data\n",
        "                # DO NOTHING\n",
        "            else:\n",
        "                print(\"[ERROR] Unexpected text data type:\" + d_type)\n",
        "                return 1\n",
        "            coordinate.append(data)\n",
        "        f.close()\n",
        "        return coordinate\n",
        "    else:\n",
        "        print(\"[ERROR] Can't read:\" + file_path)\n",
        "        return 1\n",
        "\n",
        "def drawBB(img, data):\n",
        "    color = [ COLOR_TABLE[int(data[0])][2], COLOR_TABLE[int(data[0])][1], COLOR_TABLE[int(data[0])][0] ]\n",
        "    height, width = img.shape[:2]\n",
        "    x1 = data[1]\n",
        "    y1 = data[2]\n",
        "    x2 = data[3]\n",
        "    y2 = data[4]\n",
        "    cv.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), color, 2)\n",
        "    cv.putText(img, itemIDList[int(data[0])], (x1, y1-2), cv.FONT_HERSHEY_TRIPLEX, 1, (255,255,255), 2)\n",
        "    cv.putText(img, itemIDList[int(data[0])], (x1, y1), cv.FONT_HERSHEY_TRIPLEX, 1, (0,0,0), 1)\n",
        "    return img\n",
        "\n",
        "def matching(file_list):\n",
        "    for file_path in file_list:\n",
        "        print(file_path)\n",
        "        file_name, ext = os.path.splitext(os.path.basename(file_path) )\n",
        "        file_name = file_name.replace(\"res\", \"\")\n",
        "\n",
        "        result_data = readTxt(file_path, \"result\", label_flag=LABEL_FLAG)\n",
        "        teach_data = readTxt(teach_path + file_name + \".txt\", \"teach\", cat_pass_flag=CAT_PASS_FLAG)\n",
        "        img = cv.imread(image_path + file_name + \".png\")\n",
        "        height, width = img.shape[:2]\n",
        "        img_backup = img.copy()\n",
        "\n",
        "        # Convert to Full(teach) coordinates\n",
        "        for i in range(0, len(teach_data)):\n",
        "            teach_data[i] = convNormalizedCord(teach_data[i], height, width)\n",
        "\n",
        "        # Convert for result coordinates\n",
        "        for i in range(0, len(result_data)):\n",
        "            result_data[i] = convResCord(result_data[i], height, width, NORMALIZED)\n",
        "\n",
        "        # init result lists\n",
        "        hit = [False] * len(teach_data)\n",
        "        success = [False] * len(result_data)\n",
        "        max_IoU_list = [0] * len(result_data)\n",
        "        true_category = [0] * len(result_data)\n",
        "\n",
        "        # search box which have highest IoU value\n",
        "        for j in range(0, len(result_data)):\n",
        "            max_IoU = 0.0\n",
        "            max_index = 0\n",
        "            for i in range(0, len(teach_data)):\n",
        "                iou = getIOU(teach_data[i][1:], result_data[j][1:])\n",
        "                if(max_IoU < iou) and (iou <= 1.0):\n",
        "                    max_IoU = iou\n",
        "                    max_index = i\n",
        "                if WAITTIME != 0:\n",
        "                    img = img_backup.copy()\n",
        "                    img = drawBB(img, teach_data[i])\n",
        "                    img = drawBB(img, result_data[j])\n",
        "                    cv.imshow(\"\", img)\n",
        "                    cv.waitKey(WAITTIME)\n",
        "\n",
        "            # found\n",
        "            max_IoU_list[j] = max_IoU\n",
        "            img = img_backup.copy()\n",
        "            if(max_IoU > THRESH):\n",
        "                hit[max_index] = True\n",
        "                success[j] = (teach_data[max_index][0] == result_data[j][0])\n",
        "                true_category[j] = teach_data[max_index][0]\n",
        "                if WAITTIME != 0:\n",
        "                    img = drawBB(img, teach_data[max_index])\n",
        "                    img = drawBB(img, result_data[j])\n",
        "                    cv.putText(img, \"Max IoU:\" + str(max_IoU), (0, 25) , cv.FONT_HERSHEY_TRIPLEX, 1, (255,255,255), 2)\n",
        "                    cv.putText(img, \"Class Match:\" + str(success[j]), (0, 50) , cv.FONT_HERSHEY_TRIPLEX, 1, (255,255,255), 2)\n",
        "            else:\n",
        "                if WAITTIME != 0:\n",
        "                    img = drawBB(img, result_data[j])\n",
        "                    cv.putText(img, \"No Matching\", (0, 25) , cv.FONT_HERSHEY_TRIPLEX, 1, (255,255,255), 2)\n",
        "            if WAITTIME != 0:\n",
        "                cv.imshow(\"\", img)\n",
        "                cv.waitKey(WAITTIME*5)\n",
        "\n",
        "        print(\"Result\")\n",
        "        print(\"Matching Success\")\n",
        "        for i in range(0, len(result_data)):\n",
        "            if(success[i] == True):\n",
        "                print(str(result_data[i]))\n",
        "\n",
        "        print(\"Matching Failed\")\n",
        "        for i in range(0, len(result_data)):\n",
        "            if(success[i] == False):\n",
        "                print(str(result_data[i]) + str(max_IoU_list[i]))\n",
        "\n",
        "        print(\"No Match\")\n",
        "        for i in range(0, len(teach_data)):\n",
        "            if(hit[i] == False):\n",
        "                print(str(teach_data[i]))\n",
        "\n",
        "        print(\"File output\")\n",
        "        f = open(result_path + \"matchingResults/\" + file_name + \".txt\", 'w')\n",
        "        for i in range(0, len(result_data)):\n",
        "            write_data = str(result_data[i][0]) + \" \" + str(true_category[i]) + \" \" + str(max_IoU_list[i]) + '\\n'\n",
        "            f.writelines(write_data)\n",
        "        # miss boxes\n",
        "        for i in range(0, len(teach_data)):\n",
        "            if(hit[i] == False):\n",
        "                write_data = \"0\" + \" \" + str(teach_data[i][0]) + \" 0.0\" + '\\n'\n",
        "                f.writelines(write_data)\n",
        "        f.close()\n",
        "\n",
        "    if WAITTIME != 0:\n",
        "        cv.destroyAllWindows()\n",
        "\n",
        "def evaluate(file_list):\n",
        "    conv_ID_table = [i for i in range(0, NCLASS+1)]\n",
        "\n",
        "    total_boxes = 0             # number of all(detected + not detected) boxes\n",
        "    total_true_boxes = 0        # number of true class boxes\n",
        "    total_detected_boxes = 0    # number of detected boxes\n",
        "    total_false_boxes = 0       # number of false class boxes\n",
        "    total_undetected_boxes = 0  # number of not detected boxes\n",
        "    total_IoU = 0.0\n",
        "\n",
        "    confusion_mat = [[0 for i in range(NCLASS)] for j in range(NCLASS)]\n",
        "\n",
        "    for filePath in file_list:\n",
        "        boxes = 0\n",
        "        true_boxes = 0\n",
        "        detected_boxes = 0\n",
        "        false_boxes = 0\n",
        "        undetected_boxes = 0\n",
        "\n",
        "        result_data = readTxt(filePath, \"evaluate\")\n",
        "\n",
        "        for i in range(0, len(result_data)):\n",
        "            boxes += 1\n",
        "            if(result_data[i][0] != '0' and result_data[i][1] != '0'):\n",
        "                #detected\n",
        "                detected_boxes += 1\n",
        "                total_IoU += float(result_data[i][2])\n",
        "                confusion_mat[conv_ID_table[int(result_data[i][1])]-1][conv_ID_table[int(result_data[i][0])]-1] += 1\n",
        "                if(float(result_data[i][2]) >= IOU_THRESH):\n",
        "                    #IOU >= Threshold\n",
        "                    if (result_data[i][0] == result_data[i][1]):\n",
        "                        #true class\n",
        "                        true_boxes += 1\n",
        "                    else:\n",
        "                        #false class\n",
        "                        false_boxes += 1\n",
        "                else:\n",
        "                    #low IOU\n",
        "                    false_boxes += 1\n",
        "            else:\n",
        "                #not detected\n",
        "                undetected_boxes += 1\n",
        "\n",
        "        if boxes != (detected_boxes + undetected_boxes):\n",
        "            print(\"[Error] missmatch: boxes != detected_boxes + undetected_boxes\")\n",
        "            print(boxes, detected_boxes, undetected_boxes)\n",
        "            sys.exit()\n",
        "        if detected_boxes != (true_boxes + false_boxes):\n",
        "            print(\"[Error] missmatch: detected_boxes != true_boxes + false_boxes\")\n",
        "            print(detected_boxes, true_boxes, false_boxes)\n",
        "            sys.exit()\n",
        "\n",
        "        total_boxes += boxes\n",
        "        total_true_boxes += true_boxes\n",
        "        total_detected_boxes += detected_boxes\n",
        "        total_false_boxes += false_boxes\n",
        "        total_undetected_boxes += undetected_boxes\n",
        "\n",
        "    print(\"Total Result\")\n",
        "    print(\"Matching Rate: \"+str(float(total_true_boxes) / float(total_detected_boxes)))\n",
        "    print(\"Miss(No detection box) Rate: \"+str(float(total_undetected_boxes) / float(total_boxes)))\n",
        "    print(\"Mean IoU:\"+ str(total_IoU / total_detected_boxes))\n",
        "\n",
        "    #confusion_matrix normalization\n",
        "    normalized_matrix = []\n",
        "    for i in confusion_mat:\n",
        "        a = 0\n",
        "        temp_matrix = []\n",
        "        a = sum(i,0)\n",
        "        for j in i:\n",
        "            if a == 0:\n",
        "                temp_matrix.append(0.0)\n",
        "            else:\n",
        "                temp_matrix.append(float(j) / float(a))\n",
        "        normalized_matrix.append(temp_matrix)\n",
        "\n",
        "    #draw confusion_matrix\n",
        "    plt.clf()\n",
        "    fig = plt.figure()\n",
        "    ax = fig.add_subplot(111)\n",
        "    res = ax.imshow(np.array(normalized_matrix), cmap=plt.cm.jet, interpolation='nearest')\n",
        "    cb = fig.colorbar(res)\n",
        "    cb.ax.set_yticklabels([str(i)+'%' for i in range(0, 101, 10)])\n",
        "    confusion_mat = np.array(confusion_mat)\n",
        "    width, height = confusion_mat.shape\n",
        "    item_list = [str(i) for i in range(1, NCLASS+1)]\n",
        "    plt.xticks(range(width), item_list[:width],rotation=90)\n",
        "    plt.yticks(range(height), item_list[:height])\n",
        "    plt.tick_params(labelsize=7)\n",
        "    plt.subplots_adjust(left=0.05, bottom=0.10, right=0.95, top=0.95)\n",
        "    plt.ylabel(\"True Class\")\n",
        "    plt.xlabel(\"Predicted Class\")\n",
        "\n",
        "    for i in range(0, NCLASS):\n",
        "        print(str(i+1) + \": \" + str(normalized_matrix[i][i]))\n",
        "\n",
        "    #file output\n",
        "    f = open(eval_result_path + \"/totalresult.txt\", 'w')\n",
        "    f.writelines(\"Total Result\" + '\\n')\n",
        "    f.writelines(\"Matching Rate: \"+str(float(total_true_boxes) / float(total_detected_boxes)) + '\\n')\n",
        "    f.writelines(\"Miss(No detection box) Rate: \"+str(float(total_undetected_boxes) / float(total_boxes)) + '\\n')\n",
        "    f.writelines(\"Mean IoU:\"+ str(total_IoU / total_detected_boxes) + '\\n')\n",
        "    for i in range(0, NCLASS):\n",
        "        f.writelines(str(i+1) + \": \" + str(normalized_matrix[i][i]) + '\\n')\n",
        "    f.close()\n",
        "\n",
        "    plt.savefig(eval_result_path + \"/confusion_matrix.pdf\", format=\"pdf\")\n",
        "    plt.savefig(eval_result_path + \"/confusion_matrix.png\", format=\"png\")\n",
        "\n",
        "\n",
        "def main():\n",
        "    # mkdir for Matching results\n",
        "    if not os.path.exists(result_path + \"/matchingResults\"): os.mkdir(result_path + \"/matchingResults\")\n",
        "    if not os.path.exists(result_path + \"/eval\"): os.mkdir(result_path + \"/eval\")\n",
        "\n",
        "    # IoU matching\n",
        "    file_list = glob(result_path + \"*.txt\")\n",
        "    if len(file_list) == 0:\n",
        "        print(\"[Error] Detection results file list is empty. Check this dir:\" + result_path)\n",
        "        file_list.sort()\n",
        "    else:\n",
        "        matching(file_list)\n",
        "\n",
        "    # evaluate and output\n",
        "    file_list = glob(match_result_path + \"/*.txt\")\n",
        "    if len(file_list) == 0:\n",
        "        print(\"[Error] Evaluation file list is empty. Check this dir:\" + match_result_path)\n",
        "    else:\n",
        "        file_list.sort()\n",
        "        evaluate(file_list)\n",
        "\n",
        "if __name__ == '__main__':\n",
        "    main()"
      ],
      "metadata": {
        "id": "GZKwCcHczQ52"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D1T2VOVsVpqK"
      },
      "source": [
        "Total Result以降が学習済みモデルの評価結果を表しています．\n",
        "\n",
        "Matching Rate: 0.125 \\\\\n",
        "Miss(No detection box) Rate: 0.7018633540372671 \\\\\n",
        "Mean IoU:0.7933524104872776 \\\\\n",
        "\n",
        "30epochのみの学習のため，良い結果とは言えません．本来は150epochほど学習を進める必要があります．\n",
        "\n",
        "ここまでで，SSDによる物体検出の流れを確認しました．"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3b7poEDe7UO0"
      },
      "source": [
        "## 学習済みモデルによる評価\n",
        "学習済みモデルによる評価を行います．"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "IN_DIR = \"./val/rgb\"                 # 評価データのパス\n",
        "OUT_DIR = \"./out_pretrained\"        # 結果保存先\n",
        "MODEL_PATH = \"./SSD_pretrained.pth\"  # モデルのパス\n",
        "\n",
        "# クラスラベル\n",
        "print(len(labels))\n",
        "\n",
        "def test_out2():\n",
        "    # ディレクトリの作成\n",
        "    if SAVE_DATA:\n",
        "        make_dir_list = [\n",
        "            OUT_DIR,\n",
        "            OUT_DIR + '/detection',                     # 検出結果のb-box画像\n",
        "            OUT_DIR + '/detection_txt'                  # 検出結果のb-boxテキスト\n",
        "        ]\n",
        "        for dirname in make_dir_list:\n",
        "            if not(os.path.exists(dirname)): os.mkdir(dirname)\n",
        "    # 学習モデル読み込み\n",
        "    ssd_pretrained = SSDNet()\n",
        "    ssd_pretrained.load_state_dict(torch.load(MODEL_PATH))\n",
        "    if use_cuda:\n",
        "      ssd_pretrained.cuda()\n",
        "\n",
        "    # default boxのサイズリスト計算\n",
        "    step = int(math.floor((common_params.max_ratio - common_params.min_ratio) / (len(common_params.mbox_source_layers) - 2)))\n",
        "    min_sizes = []\n",
        "    max_sizes = []\n",
        "    for ratio in range(common_params.min_ratio, common_params.max_ratio + 1, step):\n",
        "            min_sizes.append(common_params.insize * ratio / 100.)\n",
        "            max_sizes.append(common_params.insize * (ratio + step) / 100.)\n",
        "    min_sizes = [common_params.insize * 10 / 100.] + min_sizes\n",
        "    max_sizes = [common_params.insize * 20 / 100.] + max_sizes\n",
        "\n",
        "    # 入力画像の読み込み\n",
        "    color_img = glob(os.path.join(IN_DIR, '*' + '.png'))\n",
        "    color_img.sort()\n",
        "\n",
        "    # 読み込んだリストを順次検出\n",
        "    ssd_pretrained.eval()\n",
        "    for lf in range(len(color_img)):\n",
        "        img = cv.imread(color_img[lf])\n",
        "        filename, ext = os.path.splitext(os.path.basename(color_img[lf]))\n",
        "\n",
        "        if img is None:\n",
        "            print(\"[Error] 画像が読み込めません: \" + str(color_img[lf]))\n",
        "            sys.exit(1)\n",
        "        detection(img, ssd_pretrained, filename, min_sizes, max_sizes)\n",
        "\n",
        "if __name__ == '__main__':\n",
        "    test_out2()"
      ],
      "metadata": {
        "id": "fcBV_lxBzLE7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "result_path = \"./out_pretrained/detection_txt/\"\n",
        "\n",
        "# IoU matching results path\n",
        "match_result_path = result_path + \"/matchingResults/\"\n",
        "\n",
        "# evaluation output(totalresult, confusion_matrix) path\n",
        "eval_result_path = result_path + \"/eval/\"\n",
        "\n",
        "def main():\n",
        "    # mkdir for Matching results\n",
        "    if not os.path.exists(result_path + \"/matchingResults\"): os.mkdir(result_path + \"/matchingResults\")\n",
        "    if not os.path.exists(result_path + \"/eval\"): os.mkdir(result_path + \"/eval\")\n",
        "\n",
        "    # IoU matching\n",
        "    file_list = glob(result_path + \"*.txt\")\n",
        "    if len(file_list) == 0:\n",
        "        print(\"[Error] Detection results file list is empty. Check this dir:\" + result_path)\n",
        "        file_list.sort()\n",
        "    else:\n",
        "        matching(file_list)\n",
        "\n",
        "    # evaluate and output\n",
        "    file_list = glob(match_result_path + \"*.txt\")\n",
        "    if len(file_list) == 0:\n",
        "        print(\"[Error] Evaluation file list is empty. Check this dir:\" + match_result_path)\n",
        "    else:\n",
        "        file_list.sort()\n",
        "        evaluate(file_list)\n",
        "\n",
        "if __name__ == '__main__':\n",
        "    main()"
      ],
      "metadata": {
        "id": "ncN3lWw9zHtU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6F8y33QjZ_ml"
      },
      "source": [
        "Matching Rate: 0.13978494623655913\n",
        "\n",
        "Miss(No detection box) Rate: 0.14285714285714285\n",
        "\n",
        "Mean IoU: 0.8332557799526265\n",
        "\n",
        "になりました．\n",
        "\n",
        "これで，学習から評価まで一通り行うことができました．"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jbePfeZIcwZo"
      },
      "source": [
        "# 課題\n",
        "1. エポック数を変えて実験し，検出精度の変化を確認しましょう．"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#ここにコードを書く"
      ],
      "metadata": {
        "id": "ZsmcAJVDYDv3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_hT3k18MZwlH"
      },
      "source": [
        "# 参考文献\n",
        " - [1] Wei Liu, Dragomir Anguelov, Dumitru Erhan , Christian Szegedy, Scott E. Reed, Cheng-Yang Fu and Alexander C. Berg, \"SSD: Single Shot MultiBox Detector\". In 2016 European Conference on Computer Vision, pp. 21-37, 2016."
      ]
    }
  ]
}