【Albumentations】YOLO形式データセットのデータ拡張を行う

1. はじめに

Albumentationsを用いてYOLO形式データセットのデータオーギュメンテーション(Data Augmentation)を行ったため、そのやり方を備忘録のため記録しておきます。

Albumentationsとは、コンピュータビジョン用のツールで、データの水増しを簡単に行えるように用意されたライブラリです。

物体検出のタスクでは、データを水増しするときにBBoxの座標も変換する必要がありますが、そういった変換も簡単に行うことができます。
(今回はBBoxのみ解説しますが、セマンティック/インスタンスセグメンテーションのタスクで用いられるmask画像にも用いることができるようです。)

2. インストール

pipコマンドでインストールすることができます。

pip install -U albumentations

3. 実行

こちらに沿って解説していきます。

対応しているannotation形式は以下の4つのようです。
今回はyolo形式で行っていきたいと思います。

まずライブラリをインストールします。

import albumentations as A
import cv2

次にデータ拡張のパイプラインを定義します。

# 変換する関数を定義
transform = A.Compose([
    A.RandomCrop(width=400, height=400),
    A.HorizontalFlip(p=0.5),
    A.RandomBrightnessContrast(p=0.2),
], bbox_params=A.BboxParams(format='yolo',min_area=1024, 
min_visibility=0.1, label_fields=['class_labels']))

RandomCrop, HorizontalFlip, RandomBrightnessContrastは変換の関数となります。
それぞれにつきましては、以下のリンクを参照ください。

ここで1つ注意点ですが、2番目のリンク先でそれぞれの関数の説明の後ろのほうに、targetが定義されている部分があります。
このtargetにbboxesとないものは、使用することができません。

次にBboxParamsの引数について説明します。

formatはannotationの形式を指定する部分です。
YOLOの場合はformat='yolo'とします。

min_areaとmin_visibilityは、変換後のbboxの大きさが指定された閾値以下であれば、削除されるというものです。
min_areaはピクセル数で指定、min_visibilityは変化の割合で指定しています。

BBoxと別のリストでクラスラベルを与える際はlabel_fields=['class_labels']とします。

ここでBBoxと別のリストでクラスラベルを与えるとは、について説明します。
通常YOLO形式では、クラスラベル/x_center/y_center/width/heightがまとめられたファイルとなっています。

例えば、
ラベル0 / x_center=0.8 / y_center=0.7 / width=0.2 / height=0.1
ラベル1 / x_center=0.6 / y_center=0.5 / width=0.4 / height=0.3
の際、ファイルの中身としては

0 0.8 0.7 0.2 0.1
1 0.6 0.5 0.4 0.3

となります。

この場合、パイプラインに与える形式としては、
bboxes = [[0.8, 0.7, 0.2, 0.1],[0.6, 0.5, 0.4, 0.3]]
class_labels = ['0','1']
とするようにします。

これがBBoxと別のリストでクラスラベルを与えるということになります。

上記以外にも、bboxesの4つの座標の後にクラスラベルを追加して与えることも可能なようですが、上記のように分割することが好ましいと書かれていたので、今回はこちらの方法で行っていきます。

それでは画像ファイル(test.png)とアノテーションファイル(test.txt)を読み込んで、パイプラインに情報を与え、画像を表示/保存するまでのコードを示していきます。

# 画像読み込み
image = cv2.imread("test.png")
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# yoloのアノテーションファイルを読み込み
f = open('test.txt', 'r')
anno_lists = f.readlines()
f.close()

# 空の配列を用意
bboxes = []
class_labels = []

# アノテーションファイルをパイプラインに与える形に変換
for anno_list in anno_lists:
    # スペースで区切る
    anno_list = anno_list.split()

    # 0番目の要素からクラスラベルを取り出し
    class_labels.append(anno_list[0])

    # BBoxの値をfloatに変換
    anno_list = [float(i) for i in anno_list[1:]]

    # BBoxの配列に追加
    bboxes.append(anno_list)


# 変換を実行し、変換後のimage, bboxes, class_labelsを取得
transformed = transform(image=image, bboxes=bboxes, class_labels=class_labels)
transformed_image = transformed['image']
transformed_bboxes = transformed['bboxes']
transformed_class_labels = transformed['class_labels']

# RGBからBGRへ変更
transformed_image = cv2.cvtColor(transformed_image, cv2.COLOR_RGB2BGR)

# 結果を確認
cv2.imshow("img", transformed_image)
cv2.waitKey(0)
cv2.destroyAllWindows()

# 画像を書き出し
cv2.imwrite("result.png",transformed_image)

# Annotationファイルを書き出し
f = open('result.txt', 'w')

for i, transformed_bbox in enumerate(transformed_bboxes):
    f.write("{} {} {} {} {}\n".format(transformed_class_labels[i], transformed_bbox[0],
        transformed_bbox[1],transformed_bbox[2],transformed_bbox[3]))

f.close()

4.結果

変換前

変換後

画像が左右反転するのに合わせてBBoxの位置関係も反転していることがわかります。

5.まとめ

今回はAlbumentationsを用いてYOLO形式データセットのデータ拡張方法をまとめました。
YOLOv5は非常に使いやすい形で配布されているので、このデータ拡張方法と合わせることでより裾野が広がっていくのではないかと思います。