IceVisionによるObjectDetection

#これは何
2021年で最もイケている物体検出フレームワークと言っても過言ではないIceVisionを使って標準データセットのペットデータセットを使ったObject Detectionを簡単に行うためのhands-on

#IceVisionとは
過去記事参照

#環境
https://qiita.com/Ikwus/items/2a5d4b5ac2d00baf7fec

Jupyterは結局オンプレのUbuntuで動かしています。
一応google colabでの動作確認も行いました。

#手順
基本的には公式のチュートリアル「How to train the Fridge Objects Dataset」を踏襲します。
ただ、アップデートにdocument/exampleの更新が追い付いていないようなので、適宜icedataのリポジトリも参照しながら書き換えていきます。

まず、使うライブラリをimportします。

from icevision.all import *
import icedata

##モデル作成

次に、使いたいmodelを定義します。とりあえず動かすだけなのでMMDetectionのRetinanetにしておきます。
MMDetection、Torchvision、EfficientDet、Yolov5をサポートしておりmodel.(library).(model_name)のように指定することでモデル構造を定義できます。

model_type = models.mmdet.retinanet

続けてbackboneを定義します。
これもかなりサポートされており、MMDetectionに関しては対応するモデルのbackboneフォルダに使えるbackboneが一覧で記載されています。(YoLoもディレクトリ内backbone.pyにあります。)
TorchvisionとEfficientDetに関してはちょっと違うところに説明が置いてあります。開発途中だからなのかな。。。
今回はとりあえずresnet50_fpn_1xを使います。

backbone = model_type.backbones.resnet50_fpn_1x

##データセット作成

続けて、datasetを読み込みます。
今回はicevisionにインストールされている**The Oxford-IIIT Pet Dataset** (CVPR 2012)を使います。Fine-Grained Object Detectionとでも言うタスクになるんでしょうか。
ちなみに、執筆時点ではreadme.mdに記載されているfridgeデータセットは使えません。
使えるデータセットはicedataのdocumentに記載があり、
Birds (Caltech-UCSD Birds 200 Dataset)
Biwi (BIWI Sample Keypoints (center of face))
Coco (COCO Dataset)
Fridge (Fridge Objects Dataset)
Ochuman (OCHuman Dataset)
Pennfudan (Penn-Fudan Database for Pedestrian Detection and Segmentation)
Voc(PASCAL VOC 2012 Dataset)
があるようです。
ここでdatasetとclass_mapを入れておきます。

data_dir = icedata.pets.load_data()
class_map = icedata.pets.class_map()

データセットに付属のperserが icedata.(dataset).parser に標準装備されており、簡単にパーサーを設定することができます。

parser = icedata.pets.parser(data_dir = data_dir, mask=False)

このパーサーはparser.parse()を用いて分割することができますが、これにはdata_splitterのオプションがあります。

data_splitter = RandomSplitter([0.9, 0.1])
train_records, valid_records = parser.parse(data_splitter)

データセットの可視化も show_recordsを用いることで簡単にできます。

show_records(train_records[:6], ncols=3, class_map=class_map)

petsデータセットは犬派(犬:猫 = 2:1)なんですよね。

次に、データセットの形を整形します。と言ってもかなり簡単です。
このライブラリの目的は、できるだけ少ない労力でモデルを変えつつ実験を行うことなので、sizeはEfficientDetに合わせて384にしておきます。
また、augumentationはtfms.A.aug_tfmsで指定できます。(ほかのaugumentationも使えます。)
できたものをDataset()でPyTorchのDatasetクラスに変換します。

presize = 512
size = 384
train_tfms = tfms.A.Adapter([*tfms.A.aug_tfms(size=size, presize=presize), tfms.A.Normalize()])
valid_tfms = tfms.A.Adapter([*tfms.A.resize_and_pad(size), tfms.A.Normalize()])
train_ds = Dataset(train_records, train_tfms)
valid_ds = Dataset(valid_records, valid_tfms)

次にデータローダーを準備します。

train_dl = model_type.train_dl(train_ds, batch_size=16, num_workers=4, shuffle=True)
valid_dl = model_type.valid_dl(valid_ds, batch_size=16, num_workers=4, shuffle=False)

##学習

続けてモデルおよびmetricsを定義してしまいます。Retinanetの場合は以下のようになります。

model = model_type.model(backbone = backbone(pretrained = True), num_classes=len(parser.class_map))
metrics = [COCOMetric(metric_type=COCOMetricType.bbox)]

ここから先は通常のfastai (or PyTorch-Lighting)と同じ使い方になります。

learn = model_type.fastai.learner(dls=[train_dl, valid_dl], model=model, metrics=metrics)
learn.lr_find()

learn.fine_tune(20, 1e-3)

Nvidia RTX A6000で20epoch 30分程度でした。
これだけでpetの画像が学習できてしまいました。

#推論

推論も1行でできます。

model_type.show_results(model, valid_ds, detection_threshold=.5)

さすがRetinanet、そんなに長いこと計算を回してないのに、かなりいい感じに推論できています。

自前データセットでの学習だったり、指標の計算に続く...かもしれません