これまで数回に渡って、画像認識(物体検知)の主要なアルゴリズムであるSSDを改造し、精度の向上を測る方法を紹介してきました。アルゴリズムの改造が精度の向上につながっていることを確認するためには、オープンデータを使っての検証が欠かせません。物体検知用のオープンデータにはPascal VOCやCOCOがあり、前回までこれらのデータセットを使った評価を行ってきました。

Pascal VOCとCOCOのデータセットに関しては、ベースとしたSSD実装で設定を変えればそのまま利用できるよう整備されており、簡単に学習と推論を行うことができました。一方で、実際の製品やサービスに物体検知を応用する時には、目的に応じたデータを用意(入手)して学習と推測を行うことが一般的です。そこでこれから二回に渡り、カスタムデータセットを用いた学習と推測のポイントについて説明することにします。

カスタムデータセットを使う際に最初に必要となるのは、学習データを取り込むパーサーを作ることです。物体検知は学習データとして画像とアノテーションデータの組を与える必要がありますが、アノテーションデータには物体の種類(class)と位置(location)の情報が含まれます。Pascal VOCやCOCOではそれぞれ決められた形式でこれらのデータが格納されており、これを読み込むパーサーを用いて実際の学習を行います。

これまで使ってきたコードではPascal VOCとCOCOのパーサーはそれぞれ、data/voc0712.pyおよびdata/coco.pyに実装され、それぞれに含まれるVOCDetectionクラスやCOCODetectionクラスがパーサーの本体です。カスタムデータセットに対応するには、実際のデータのフォーマットに合わせ、これらのクラスを実装すれば良いことになります。

実装するクラスはtorchvision.datasets.vision.VisionDatasetを継承し、最低限__init__、__getitem__、__len__のメソッドを実装する必要があります。__init__はクラスインスタンスの初期設定、__len__は文字通りデータセット長を返すメソッドですので、パーサーの本体はデータ要求に応じる__getitem__が担います。

data/voc0712.pyの__getitem__の実装を見ると以下のようになっており、最終的にデータセット中のindexで指定されたデータのイメージデータとアノテーションデータのセットを返すと良さそうです。

    def __getitem__(self, index):
        im, gt, h, w = self.pull_item(index)
       return im, gt

返り値となるイメージデータ(im)とアノテーションデータ(gt)のフォーマットですが、イメージデータはchannel firstでTorch Tensorフォーマットとする必要があります。RGBの画像データはChannel(RGB)、H(縦)、W(横)の三次元データで表現されますが、この次元の順序がPytorchでは CHWとなり、OpenCVなどで読み込んだchannel last(HWC)のデータからは変換する必要があるのです。

また、アノテーションデータは、物体クラスを数値に変換し、ロケーションデータは画像全体の幅に対する比率で、[xmin, ymin, xmax, ymax, class]の情報として与えます。

最終的に一つの画像に含まれるオブジェクトの情報はこのような形にまとめられます。 
[[0.03892544 0.19088146 0.12335526 0.27234043 2.        ]
 [0.31524123 0.19574468 0.32620614 0.21702128 4.        ]
 [0.34868421 0.18966565 0.36513158 0.22796353 4.        ]
 [0.09375    0.1775076  0.13377193 0.23221884 2.        ]
 [0.07236842 0.15379939 0.08497807 0.16352584 5.        ]
 [0.15789474 0.14650456 0.16995614 0.15744681 5.        ]
 [0.36074561 0.10151976 0.63541667 0.32401216 9.        ]
 [0.21326754 0.18966565 0.22039474 0.20243161 4.        ]]

アノテーションデータが以下のようなJSONフォーマットで与えられた時のパーサーを考えてみます。
{
    "labels": [
        {
            "box2d": {
                "x1": 1651,
                "x2": 1827,
                "y1": 584,
                "y2": 676
            },
            "category": "Truck"
        },
        {
            "box2d": {
                "x1": 1793,
                "x2": 1934,
                "y1": 581,
                "y2": 682
            },
            "category": "Truck"
        }
    ]
}

以下はこの形式のデータセットを読み込むパーサーの例です。
import torch
import torchvision

import os
import cv2
import json
import numpy as np

class JMEDetection(torchvision.datasets.vision.VisionDataset):

    CLASSES = ["Bicycle",
               "Bus",
               "Car",
               "Motorbike",
               "Pedestrian",
               "Signal",
               "Signs",
               "SVehicle",
               "Train",
               "Truck"]
    
    def __init__(self,
                 root,
                 transform=None,
                 target_transform=None,
                 transforms=None):
        super(JMEDetection, self).__init__(root, transforms, transform, target_transform)
        self.name = 'JME'

        jme_root = '/mnt/ssd/signate/'
        image_dir = jme_root   'dtc_train_images/'
        annotation_dir = jme_root   'dtc_train_annotations/'

        filenames = jme_root   'train.txt'
        
        # filename list
        with open(filenames, "r") as f:
            file_names = [x.strip() for x in f.readlines()]

            

        self.images = [os.path.join(image_dir, x   ".jpg") for x in file_names]
        self.annotations = [os.path.join(annotation_dir, x   ".json") for x in file_names]

        assert (len(self.images) == len(self.annotations))

    def __getitem__(self, index):
        """
        Args:
            index (int): Index

        Returns:
            tuple: (image, target) where target is a dictionary.
        """
        img = cv2.imread(self.images[index])
        h, w, _ = img.shape

        with open(self.annotations[index]) as f:
            info = json.load(f)
            boxes = np.empty((0,4))
            labels = np.array([])
            for label in info['labels']:
                xmin = int(label['box2d']['x1'])/w
                xmax = label['box2d']['x2']/w
                ymin = label['box2d']['y1']/h
                ymax = label['box2d']['y2']/h
                cls = JMEDetection.CLASSES.index(label['category'])
                boxes = np.append(boxes, [[xmin, ymin, xmax, ymax]], axis=0)
                labels = np.append(labels, cls)

        if self.transform is not None:
            img, boxes, labels = self.transform(img, boxes, labels)
            img = img[:, :, (2, 1, 0)]  # BGR -> RGB
            target = np.hstack((boxes, np.expand_dims(labels, axis=1)))
            # H,W,C -> C,H,W
        return torch.from_numpy(img).permute(2, 0, 1), target

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

実際に学習を行う際にはdata augmentationと呼ばれるデータの前処理が行われます。data augmentationでは左右をランダムに入れ替えるrandom flipや画像の一部をrandomに取り出すrandom cropを行って、データセットの拡張を行っています。(dataからmeanを引いたりsdで割ったりする正規化を行うこともあります) 

SSDで行われる標準的なaugmentationはutils/augmentations.pyにまとめられていますのでここから選んで使用することができます。上記の例ではクラスの初期化を行う際にtransformの値として、augmentationのインスタンスを与えます。

    dataset = JMEDetection(root=args.dataset_root,  transform=SSDAugmentation(cfg['size'], MEANS))

Pascal VOCでは300x300あるいは512x512の入力解像度を想定してモデルを設計していました。実際のカスタムデータセットを使うときにはもっと高い解像度のデータが得られることが多く、また画像の縦横比が1:1でないことがほとんどです。データセットの解像度やアスペクト比にあわせモデルとアンカーボックスを設計ことが、カスタムデータセットを利用して物体検知を行う残りの手順となります。

次回はその部分を説明するとともに、カスタムデータセットを用いた物体検知の学習および推測の全体のコードを紹介することにします。