【第二回】YOLOv3を一から実装 ～dataset作成～

はじめに

　前回の記事ではYOLOv3のモデル構造について簡単に解説しました。今回は実際にYOLOv3を実装するにあたって、動作の確認等に使うためにdatasetを先に作成していこうと思います。データは実際に論文でも使用されたCOCO datasetを使ってdatasetを作成していきます。

COCO datasetは物体検出系の論文でよく評価用に使われています。本記事ではCOCO datasetを扱うCOCO API(pycocotools)を使って、pytorchのdatasetを作るまでを解説します。

最終目標

　最終的にdatasetにindexを入力するとpixel値を0-1に正規化した画像([color channel, height, width]の順)、画像内の物体の位置を示すバウンディングボックスの座標[x_center, y_center, height, width]、バウンディングボックスの内の物体のラベル(種類)を含むTensorが返ってくるようなdatasetを作っていきます。図で示すと以下のようなイメージです。

画像によってバウンディングボックスの数は異なるので取得するボックスの最大数を指定します。指定した最大数より少ない数のバウンディングボックスを持つ場合は逆に最大数に揃えるようにボックスの座標をすべて0として取得するようにしています。

フォルダ構成

作業フォルダ
　┠ COCO (COCO datasetのダウンロードの項で作成します)
　┗ dataset
　　　┗ cocodataset.py (今回作成するファイル)

使用ライブラリ

　使用ライブラリは以下となります。Google Colabを使用する場合はpycocotoolsも含めすべてインストールされているのでpipは不要です。

cocodataset.py

import os
import cv2
import numpy as np
from pycocotools.coco import COCO

import torch
from torch.utils.data import Dataset

COCO datasetのダウンロード

　COCO APIはあくまでローカルのフォルダ内にあるCOCO datasetを読み取るAPIとなっているので、事前にデータをダウンロードする必要があります。ここではlinuxのコマンドでダウンロードしていきましょう。linuxコマンドもColabやJupter Notebookのセルに打ち込んで実行することができます。

始めにCOCOというディレクトリを作成し、カレントディレクトリをCOCOに指定します。

!mkdir COCO
%cd COCO

続いて画像データとannotationデータをダウンロードします。2014年版と2017年版があるのですが今回はCOCO dataset 2017のvalidataionデータをダウンロードします。trainデータもあるのですが画像の方がzip形式で18GBの容量なのでvalidationデータでいきます。

%%bash
wget http://images.cocodataset.org/zips/val2017.zip
wget http://images.cocodataset.org/annotations/annotations_trainval2017.zip

以上でCOCOというフォルダ内にval2017.zipとannotations_trainval2017.zipがダウンロードできているはずです。

　最後にzipファイルを展開して、ついでにzipファイルの方は使わないので削除してしまいましょう。

%%bash
unzip val2017.zip
unzip annotations_trainval2017.zip

rm -f val2017.zip
rm -f annotations_trainval2017.zip

COCODatasetのコンストラクタ

　始めにCOCODatasetのコンストラクタを作成していきます。pytorchのDatasetクラスを継承しています。

cocodataset.py

class COCODataset(Dataset):
    def __init__(self, img_size=416, min_size=1, max_labels=10):
        # ① パス指定
        self.data_dir = './COCO/'
        self.json_file = 'instances_val2017.json'
        self.img_dir = 'val2017'

        # ② instances_val2017.json解析
        self.coco = COCO(self.data_dir + 'annotations/' + self.json_file)

        # ③ 画像idとクラスidを取得
        self.ids = self.coco.getImgIds()
        self.class_ids = sorted(self.coco.getCatIds())

        # ④ その他初期設定
        self.img_size = img_size
        self.min_size = min_size
        self.max_labels = max_labels

①まず各パスを指定します。②続いてinstances_val2017.jsonファイルまでのパスをCOCO()メソッドに入力することでjsonファイルを解析してくれます。

③解析したself.cocoに対してgetImgIds()メソッドを使うことですべての画像のidを取得、getCatIds()メソッドを使うことで分類するクラスのidを取得することができます。COCO2017は80種類のラベルが各位置情報につけられていますが中を試しに見ていると...

cocodataset.py

dataset = COCODataset()
print(len(dataset.class_ids))
print(dataset.class_ids)

"""
80
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90]
"""

80種類クラスがありますがクラスのidは連番ではないみたいですね。

最後に④その他の設定ですが、img_sizeは画像のサイズの指定で、後の処理で画像を指定した値の正方形にリサイズするようにします。min_sizeとmax_labelsも後の処理で使うのですがその時に説明します。

dataset内の画像数を取得するメソッド実装

　datasetに含まれる全画像の枚数を取得したい場合が出てくるので、ここで特殊メソッド __len__ で枚数を取得できるようにしておきます。画像枚数は画像のidの数を調べることで取得できます。

cocodataset.py

class COCODataset(Dataset):
    def __init__(self, img_size=416, min_size=1, max_labels=10):
      　# 省略
　　　
　　def __len__(self):
        return len(self.ids)

確認してみましょう。

cocodataset.py

dataset = COCODataset()
print(len(dataset))

"""
5000
"""

COCOval2017には5000枚の画像が含まれているようですね。

画像とラベルを取得するメソッドの実装

　indexを入力してその画像やラベルの情報を取得するために特殊メソッド __getitem__ を実装していきます。まずは画像を取得する部分から作成していきましょう。

cocodataset.py

class COCODataset(Dataset):
    def __getitem__(self, index):
        id_ = self.ids[index]
        
        # load image
        img_file = os.path.join(self.data_dir, self.img_dir, '{:012}'.format(id_) + '.jpg') # 12桁になるよう0パディング
        img = cv2.imread(img_file)

引数でindexを受け取り、その引数番目の画像idを取得します。続いて画像のidから画像のパスを作成します。画像のファイル名はidを右詰めにして合計12桁になる用0でパディングされています(例:000000037777.jpg)。'{:012}'.format(id_) が0パディングしている部分に当たります。最後にcv2で画像ファイルを読み込みます。

　YOLOv3に入力する画像は縦横同じpx値の正方形である必要があるのでリサイズする必要があります。しかし単純にcv2のresizeメソッドでresizeを行うと下の図の左下のように画像の縦横比まで変わってしまいます。YOLOv3ではこのようなリサイズではなく、図右下のように比率を変えずにリサイズする方が精度が良くなるとされているため右下のように余白をグレーに埋めるようなリサイズをする関数を別途用意していきます。

　縦横の比を維持したまま指定のサイズに画像をリサイズする関数が以下です。

cocodataset.py

def pad_resize(img, img_size: int):
    h, w, _ = img.shape
    img = img[:, :, ::-1]   # BGR >> RGB

    r = w / h               # 比率

    ## 縦長の場合
    if r < 1:
        new_h = img_size        # new_h = imgsize
        new_w = new_h * r       # new_w = new_h * w / h << h:w = new_h : new_w
    ## 横長の場合
    else:
        new_w = img_size
        new_h = new_w / r 

    new_w, new_h = int(new_w), int(new_h)

    # 拡大する量(片側)
    dx = (img_size - new_w) // 2
    dy = (img_size - new_h) // 2

    # resize
    img = cv2.resize(img, (new_w, new_h))                           # 一旦new_wとnew_hでリサイズ

    sized = np.ones((img_size, img_size, 3), dtype=np.uint8) * 127
    sized[dy:dy+new_h, dx:dx+new_w, :] = img                        # 余白を127で埋めて入れる

    img_info = (h, w, new_h, new_w, dx, dy)
    return sized, img_info

　コードの詳細は中身を見ていただければと思います。この関数にcv2で読み込んだ画像を入力すると以下変換が行われます。

・カラーチャンネルの順番をBGRからRGBに変換
・画像サイズを比率を変えずリサイズし、余白部分はpixel値127で埋める

関数の出力は上記変換が行われた画像(ndarray型)に加えて、「元画像の幅、元画像の高さ、リサイズ後の画像部分の幅、リサイズ後の画像部分の高さ、127で埋めた余白の幅(片側)、127で埋めた余白の高さ(片側)」の6要素の情報を含んだimg_info(タプル型)を返します。画像のリサイズと一緒にバウンディングボックスもリサイズする必要があり、その際に使用する情報となっています。

　それでは元のCOCOdatasetクラスの作成に戻ります。読み込んだ画像を先ほど作成したpad_resize()関数でリサイズし、さらにpixel値が0-1に収まるように正規化します。

cocodataset.py

class COCODataset(Dataset):
    def __init__(self, img_size=416, min_size=1, max_labels=10):
      　# 省略
　　　
　　def __len__(self):
        return len(self.ids)

　　def __getitem__(self, index):
        id_ = self.ids[index]
        
        # load image
        img_file = os.path.join(self.data_dir, self.img_dir, '{:012}'.format(id_) + '.jpg') # 12桁になるよう0パディング
        img = cv2.imread(img_file)
        img, img_info = pad_resize(img, self.img_size)
        img = np.transpose(img / 255, (2, 0, 1))                                            # px値0-1化 + [ch, h, w]に並び替え

　さて次はラベル部分の読み込みについて説明していきます。annotationデータのidはgetAnnIds()メソッドで取得します。引数imgIdsに画像のidを指定することで対応したannotationデータの取得ができます。さらに取得したannotationデータのidをloadAnns()メソッドに入力することで指定した画像に含まれるクラスラベルや位置情報を含むannotationデータが得られます。

cocodataset.py

class COCODataset(Dataset):
    def __init__(self, img_size=416, min_size=1, max_labels=10):
      　# 省略
　　　
　　def __len__(self):
        return len(self.ids)

　　def __getitem__(self, index):
        id_ = self.ids[index]

        # load image
        img_file = os.path.join(self.data_dir, self.img_dir, '{:012}'.format(id_) + '.jpg')     # 12桁になるよう0パディング
        img = cv2.imread(img_file)
        img, info_img = pad_resize(img, self.img_size)
        img = np.transpose(img / 255., (2, 0, 1))       # px値0-1化 + [ch, h, w]に並び替え
       
        # load labels
        anno_ids = self.coco.getAnnIds(imgIds=[int(id_)], iscrowd=None)
        annotations = self.coco.loadAnns(anno_ids)

取得したバウンディングボックスとラベル情報を含むannotationsをforループでlabelsというリストに格納していきます。その際にインスタンス化時に指定するmin_sizeの値より小さい高さ、幅のバウンディングボックスであるラベルは無視するようにしています。

cocodataset.py

class COCODataset(Dataset):
    def __init__(self, img_size=416, min_size=1, max_labels=10):
      　# 省略
　　　
　　def __len__(self):
        return len(self.ids)

　　def __getitem__(self, index):
        # 省略
       
        # load labels
        anno_ids = self.coco.getAnnIds(imgIds=[int(id_)], iscrowd=None)
        annotations = self.coco.loadAnns(anno_ids)

        for anno in annotations:
            # h, w が min_sizeを超えているか場合labelsに追加
            if (anno['bbox'][2] > self.min_size) and (anno['bbox'][3] > self.min_size):
                labels.append([])                                               # labelsに空のリスト追加
                labels[-1].append(self.class_ids.index(anno['category_id']))    # 追加したリストの最後尾にcategory_idを格納
                labels[-1].extend(anno['bbox'])                                 # 追加したリストの最後尾にbboxを格納

　ここで新たな関数を準備します。COCOデータセットで用意されているバウンディングボックス(BB)の位置情報は[BB左上x座標, BB左上y座標, height, width]で構成されています。一方YOLOv3で使用する位置情報は[BB中心x座標, BB中心y座標, height, width]であるので、変換する関数label2yolobox()関数を用意します。

cocodataset.py

def label2yolobox(labels, img_info, maxsize):
    h, w, nh, nw, dx, dy = img_info

    x1 = labels[:, 1] / w
    y1 = labels[:, 2] / h
    x2 = (labels[:, 1] + labels[:, 3]) / w
    y2 = (labels[:, 2] + labels[:, 4]) / h

    labels[:, 1] = (((x1 + x2) / 2) * nw + dx) / maxsize
    labels[:, 2] = (((y1 + y2) / 2) * nh + dy) / maxsize
    labels[:, 3] *= nw / w / maxsize
    labels[:, 4] *= nh / h / maxsize

    return labels

関数にはラベルに加えて、pad_resize()関数でリサイズしたときの出力img_infoとリサイズ後の画像サイズmaxsizeを入力します。すると出力としてリサイズ後の画像サイズに合わせた、YOLO用のバウンディングボックスが出力として得られるようになっています。

　作成したlabel2yolobox()関数を使ってラベルの変換と、ラベルの数がmax_labelsを超えている、または満たしていない場合にmax_labelsに数を合わせるようにする処理を追加し、出力まで仕上げてしまいます。

cocodataset.py

class COCODataset(Dataset):
    def __init__(self, img_size=416, min_size=1, max_labels=10):
      　# 省略
　　　
　　def __len__(self):
        return len(self.ids)

　　def __getitem__(self, index):
        # 省略
       
        # load labels
        anno_ids = self.coco.getAnnIds(imgIds=[int(id_)], iscrowd=None)
        annotations = self.coco.loadAnns(anno_ids)

        for anno in annotations:
            # h, w が min_sizeを超えているか場合labelsに追加
            if (anno['bbox'][2] > self.min_size) and (anno['bbox'][3] > self.min_size):
                labels.append([])                                               # labelsに空のリスト追加
                labels[-1].append(self.class_ids.index(anno['category_id']))    # 追加したリストの最後尾にcategory_idを格納
                labels[-1].extend(anno['bbox'])                                 # 追加したリストの最後尾にbboxを格納

            padded_labels = np.zeros((self.max_labels, 5))                       # max_labels × 5 (class, x, y, h, w)

            # ラベルが見つかったら
            if len(labels) > 0:
                labels = np.stack(labels)                                       # list to ndarray
                labels = label2yolobox(labels, img_info, self.img_size)
                padded_labels[range(len(labels))[:self.max_labels]] = labels[:self.max_labels]

            # label ndarray to tensor
            padded_labels = torch.from_numpy(padded_labels)

            return img, padded_labels

　まずpadded_labelsというmax_labels × 5 のサイズで要素がすべて0の配列を用意します。もし画像にmin_sizeを超えるラベルが一つもない場合はこのpadded_labelsが画像のラベルとしてそのまま出力されます。ラベルが存在する場合はif文内の処理が行われます。具体的にはラベルをlabels2yolobox()関数で変換し、padded_labelsに上書きします。

　以上でdatasetの作成は完了です。動作確認をしていきます。

cocodataset.py

dataset = COCODataset()

img, labels = dataset[8]
print(img.shape)
print(labels.size())

'''
(3, 416, 416)
torch.Size([10, 5])
'''

COCODatasetをインスタンス化しindexが8の画像とラベルを取得すると、[3, 416, 416]の画像と10個のラベル情報が返ってきていることが分かります。

まとめ

　ここまで読んでいただきありがとうございます。COCO APIを使ってYOLOv3用にカスタマイズdatasetの作成を説明しました。次回はYOLOv3のモデル構造部分を解説したいと思います。最後に今回のコードの全体を載せておきます。

コードまとめ

cocodataset.py

import os
import cv2
import numpy as np
from pycocotools.coco import COCO

import torch
from torch.utils.data import Dataset

# input     >> img(cv2で読み込んだndarray), imasize(リサイズ後のサイズ)
# output    >> img(ndarray), info_img([元のh, 元のw, リサイズ後画像が存在するh, リサイズ後画像が存在するw, 127で埋めた横方向の片側量, 127で埋めた縦方向の片側量])
def pad_resize(img, img_size: int):
    h, w, _ = img.shape
    img = img[:, :, ::-1]   # BGR >> RGB

    r = w / h               # 比率

    ## 縦長の場合
    if r < 1:
        new_h = img_size        # new_h = imgsize
        new_w = new_h * r       # new_w = new_h * w / h << h:w = new_h : new_w
    ## 横長の場合
    else:
        new_w = img_size
        new_h = new_w / r 

    new_w, new_h = int(new_w), int(new_h)

    # 拡大する量(片側)
    dx = (img_size - new_w) // 2
    dy = (img_size - new_h) // 2

    # resize
    img = cv2.resize(img, (new_w, new_h))                           # 一旦new_wとnew_hでリサイズ

    sized = np.ones((img_size, img_size, 3), dtype=np.uint8) * 127
    sized[dy:dy+new_h, dx:dx+new_w, :] = img                        # 余白を127で埋めて入れる

    img_info = (h, w, new_h, new_w, dx, dy)
    return sized, img_info

# input     >> labls: [x(左上), y(左上), h, w]
# output    >> labes: [xcenter, ycenter, h, w] (pad_resizeに合わせて変換)
def label2yolobox(labels, img_info, maxsize):
    h, w, nh, nw, dx, dy = img_info

    x1 = labels[:, 1] / w
    y1 = labels[:, 2] / h
    x2 = (labels[:, 1] + labels[:, 3]) / w
    y2 = (labels[:, 2] + labels[:, 4]) / h

    labels[:, 1] = (((x1 + x2) / 2) * nw + dx) / maxsize
    labels[:, 2] = (((y1 + y2) / 2) * nh + dy) / maxsize
    labels[:, 3] *= nw / w / maxsize
    labels[:, 4] *= nh / h / maxsize

    return labels  


class COCODataset(Dataset):
    def __init__(self, img_size=416, min_size=1, max_labels=10):
        # ① パス指定
        self.data_dir = './COCO/'
        self.json_file = 'instances_val2017.json'
        self.img_dir = 'val2017'

        # ② instances_val2017.json解析
        self.coco = COCO(self.data_dir + 'annotations/' + self.json_file)

        # ③ 画像idとクラスidを取得
        self.ids = self.coco.getImgIds()
        self.class_ids = sorted(self.coco.getCatIds())

        # ④ その他初期設定
        self.img_size = img_size
        self.min_size = min_size
        self.max_labels = max_labels

    def __len__(self):
        return len(self.ids)

    def __getitem__(self, index):
        id_ = self.ids[index]
        
        # load image
        img_file = os.path.join(self.data_dir, self.img_dir, '{:012}'.format(id_) + '.jpg') # 12桁になるよう0パディング
        img = cv2.imread(img_file)
        img, img_info = pad_resize(img, self.img_size)
        img = np.transpose(img / 255, (2, 0, 1))                                            # px値0-1化 + [ch, h, w]に並び替え

        # load labels
        anno_ids = self.coco.getAnnIds(imgIds=[int(id_)], iscrowd=None)
        annotations = self.coco.loadAnns(anno_ids)

        labels = []

        for anno in annotations:
            # anno: [x(左上), y(左上), h, w]

            # h, w が min_sizeを超えているか場合labelsに追加
            if (anno['bbox'][2] > self.min_size) and (anno['bbox'][3] > self.min_size):
                labels.append([])                                               # labelsに空のリスト追加
                labels[-1].append(self.class_ids.index(anno['category_id']))    # 追加したリストの最後尾にcategory_idを格納
                labels[-1].extend(anno['bbox'])                                 # 追加したリストの最後尾にbboxを格納


            padded_labels = np.zeros((self.max_labels, 5))                       # max_labels × 5 (class, x, y, h, w)

            # ラベルが見つかったら
            if len(labels) > 0:
                labels = np.stack(labels)                                       # list to ndarray
                labels = label2yolobox(labels, img_info, self.img_size)
                padded_labels[range(len(labels))[:self.max_labels]] = labels[:self.max_labels]

            # label ndarray to tensor
            padded_labels = torch.from_numpy(padded_labels)

            return img, padded_labels, img_info, id_