技術は使ってなんぼ

自分が得たものを誰かの役に立てたい

【SSD】物体検出にCutMixを実装して、精度向上するか検証してみた

背景

YOLO v4でも高い精度向上に貢献したとされる新しい技術の一つに「CutMix」というものがある。


この技術がどのぐらい精度向上に貢献するのか、検証してみたい。

CutMixとは?

augmentation(データ拡張)の技術の一つで、MixUpとCutoutという従来のaugmentationの欠点を改善する手法です。


CutMixの論文や参考にしたGit
github.com


分かりやすい絵がこれ。上記Git内の論文から参照しています。

f:id:yonesuke0716:20201008233650p:plain
CutMix

やってることもそんなに難しい手法ではないため、実装も比較的容易と想像します。


論文もとてもシンプルで読みやすかったです。


ただし、論文中にも記載がありますが、物体検出にはそのまま実装することが難しく、本論文でも転移学習による実装にとどめているようです。

どういうこと?

この論文の実装は、クラス分類での実装となります。


クラス分類の教師データは、画像データ一枚に対してone-hotベクトルのような正解label(大体list型)が与えられます。


一方物体検出では、画像データ一枚に対して複数のbboxと、それぞれに対応した正解label(大体ID等のint型)が与えられます。


特にCutoutのように、画像の一部を切り取るため、bboxも不連続な形で切り取られてしまう可能性があります。


難しいのは、その切り取られたbboxと対応するlabelをどう処理するかということです。

どうやったか?

CutMixのやりたいことの要点は、
1.CutOutのような、元データの一部を切り出すこと
2.その切り取った部分を効率よく使うため、別のデータを入れること


これにより効率よく効果的な学習が実現できていると考えられます。


なので私は、
1.切り取る部分はbboxひとつをそのまま切り取って貼り付ける
2.張り付けた先のbboxとlabelに、切り取ったデータをそのままくっつける

とすることで、上記のbboxが不連続になることなく、Cutoutの効果を得ることができると考えました。


ただし、この方法だとbboxの大きさがあまりに小さいと大して効果がでないかもというデメリットもあります。


まぁその辺はbboxのサイズが閾値以上に限定したり、複数のbboxで行う等の改良をすれば解決するかと思いますが、一旦上記の方法で実装して効果を検証します。

ではさっそく試してみましょう!

今回は「Blood_Cell_Detection」というSSDを使った学習モデルで効果を確認していきたいと思います。


Google Colabで既に動くnotebookがあったので、こちらを拝借いたしました。


SSD等の物体検出に興味がある方は、こちらを動かすだけでも勉強になりますよ。

colab.research.google.com


こちらのデータセットは物体検出のデータセットとしてはかなり数が少ない軽量タイプなので、比較的短時間で検証ができます。


クラスとしては赤血球(Red Blood Cell; RBC)、白血球(White Blood Cell; WBC)、血小板(Platelet)の3クラスで、各bbox毎に振られています。


試しに2枚使って、CutMixやってみます。

f:id:yonesuke0716:20201008235313p:plain
切り取り元画像データ
f:id:yonesuke0716:20201008235341p:plain
貼り付け先画像データ


この2枚をCutMixしたのがこちら。どこが変わったか、わかりますか?

f:id:yonesuke0716:20201008235453p:plain
CutMix画像


真ん中の右下のbboxが合成されました。

ソースコード

上記Colabの「Data augmentationの実装」のdef __call__を以下のように変更しました。

def __call__(self, in_data):
# There are five data augmentation steps
# 0. CutMix
# 1. Color augmentation
# 2. Random expansion
# 3. Random cropping
# 4. Resizing with random interpolation
# 5. Random horizontal flipping

    img, bbox, label = in_data
    # 0. CutMix
    if np.random.randint(1,7) % 2 == 0:
        idx = random.randint(0, len(train_dataset)-1)
        cut_img, bbox_list, label_list =  train_dataset[idx]
        bb_label_id = random.randint(0, len(bbox_list)-1)
        cut_bbox = bbox_list[bb_label_id]
        cut_label = label_list[bb_label_id]

        cut_img = cut_img.astype(np.int32)
        cut_bbox = cut_bbox.astype(np.int32)
        cut_label = cut_label.astype(np.int32)

        def cutmix(img_1, img_2, bbox_1, bbox_2, label_1, label_2):
            by1, bx1, by2, bx2 = bbox_2
            img_1[:, by1:by2, bx1:bx2] = img_2[:, by1:by2, bx1:bx2]
            new_label = np.append(label_1, label_2)
            new_bbox = np.append(bbox_1, [list(bbox_2)], axis=0)
            return img_1, new_bbox, new_label

        img, bbox, label = cutmix(img, cut_img, bbox, cut_bbox, label, cut_label)
        #ax = vis_bbox(img, bbox, label, label_names=bccd_labels)
        #ax.set_axis_off()
        #ax.figure.tight_layout()

    # 1. Color augmentation
    img = random_distort(img)

解説

まずCutMixを実施する場所を先頭にしました。


もし後ろの方にもっていくと、途中で別のAugmentationによりbboxの位置が大きく変化してしまうため、合成時にエラーをはいてしまうことがわかりました。


Color augmentationの下でも大丈夫と思います。

if np.random.randint(1,7) % 2 == 0:

Augmentationも毎回やればいいってもんでもないので、とりあえずサイコロの確率で更に偶数の時、つまり50%で実施するようにしています。


この辺のランダム確率は好きなように調整していいと思います。

idx = random.randint(0, len(train_dataset)-1)
cut_img, bbox_list, label_list =  train_dataset[idx]
bb_label_id = random.randint(0, len(bbox_list)-1)
cut_bbox = bbox_list[bb_label_id]
cut_label = label_list[bb_label_id]

もう一枚画像データをランダムでとってきます。とってきた画像データに付随するbboxは2次元のnumpyなので、そこから更にランダムで一個取得し、同じ正解ラベルも一つ取得します。

cut_img = cut_img.astype(np.int32)
cut_bbox = cut_bbox.astype(np.int32)
cut_label = cut_label.astype(np.int32)

これは後ろで出てくる、img_1[:, bx1:bx2, by1:by2] = img_2[:, bx1:bx2, by1:by2]の合成対策です。

float型は受け付けないため、int型に変更してます。

def cutmix(img_1, img_2, bbox_1, bbox_2, label_1, label_2):
    bx1, by1, bx2, by2 = bbox_2
    img_1[:, bx1:bx2, by1:by2] = img_2[:, bx1:bx2, by1:by2]
    new_label = np.append(label_1, label_2)
    new_bbox = np.append(bbox_1, [list(bbox_2)], axis=0)
    return img_1, new_bbox, new_label

CutMixの肝となる部分です。今回のデータセットのchainercvによる画像データは「C, H, W」の順なので「img_1[:, bx1:bx2, by1:by2] = img_2[:, bx1:bx2, by1:by2]」という並びにすることで、bbox部分の合成が可能となります。


後は、貼り付け先のlabelとbboxに追加すれば、やりたいことが実現できます。

学習結果

さて一番の問題は、論文とは異なる「なんちゃってCutMix」で精度が向上するのか?というところです。


実際に学習を回してみましょう。


このColabでは、エポック数300回のうち既に290回回してくれているモデルが用意されています。


残りの10回を転移学習という形で実施するようです。


まぁわずか10回ではそんなに変わらないだろうと思いながらも、いざ検証!

f:id:yonesuke0716:20201009002356p:plain
評価結果


お、ちょっと上がってる!赤血球(Red Blood Cell; RBC)が0.1%上がってます!逆に血小板(Platelet)は0.0015%下がってますねぇ。これぐらいまでくると誤差かな?


まぁ、わずか10回の学習なので、たまたま上がったのかもしれませんが、それでも0.1%のインパクトは大きい!


CutMix、なかなか期待できるかも。もう少しCutするbboxの数増やすかなぁ。


epoch300回で回してみたのですが、残念ながらgoogle colabのメモリ12GBではメモリーが足りなくて途中で止まってしまいました。。


まだまだ色々検証する余地がありそうですが、今回はここまで!