物体検出のためのアノテーション作業・データ拡張

機械学習や深層学習を学び始めて2か月の者です。
Yolo v3による物体検出を行う際に、準備作業としてアノテーションやデータ拡張を行う必要があったのですが、自分がやりたかったことや自分の使用している環境に完全に合致した記事が見当たらなかったので、せっかくなので自分が取り組んだ内容を備忘録として残すことにしました。
より良い案がありましたらご意見・アドバイス頂けましたら幸いです。

はじめに

物体検出の基本事項について簡単に触れておきます。

物体検出とは
写真や動画データの中から犬や猫などの物体を検出し、検出した物体の周囲を四角い箱（バウンディングボックス）で囲う深層学習の手法のことです。

アノテーションとは
物体検出を行うための深層学習モデルを作るに当たり、「これは犬だよ」「これは猫だよ」といった感じで画像データ内のオブジェクトの周りをバウンディングボックスで囲ってやり、それをモデルにインプットして学習させる必要があります。アノテーションとは、そのバウンディングボックスを手作業で画像に付与する作業のことを指します。

データ拡張とは
深層学習モデルの学習の際に使用する教師データが少ない場合に、データの水増しのために画像を少し拡大/縮小したり、あるいは回転したりするなどの加工を行う作業のことを指します。

やったこと

主に下記の作業になります。

【作業1】アノテーション作業
これには「labelImg」というツールを使用しました。このツールを使ってjpg等の画像ファイルに手作業でバウンディングボックスを囲ってやることで、そのボックスの座標情報を含んだXMLファイルが生成されます。今回私が使用した環境では、「Pascal VOC」というフォーマットで作成しました。
(使用方法 参考：https://note.com/npaka/n/nf74e32b47712)

【作業2】アノテーションで作成したXMLファイルの取り込み
これには、Pythonに組み込まれている「xml.etree.ElementTree」というXML読み込み用のライブラリを用いて、XMLファイル内のバウンディングボックスの情報をPythonプログラムに読み込ませました。
読み込むための処理は、下記サイトのサンプルコードを参考にさせて頂きつつ、自分がやりたいように少しカスタマイズしました。
(参考：http://maruo51.com/2020/06/05/pascalvoc_dataset/)

【作業3】データ拡張作業
データ拡張処理に使用できるライブラリとしては、TensorFlowに備わっている「ImageDataGenerator」が有名かと思いますが、ここでは上記のアノテーション作業によって画像に付与されたバウンディングボックスの情報ごとデータ拡張したかったため、代わりに「Albumentations」というライブラリを使用しました。
【作業2】でXMLファイルから取り込んだバウンディングボックスの座標情報とオリジナルの画像データをこのAlbumentationsのAPIにインプットすることで、画像データに拡大/縮小等の加工を加えた際にそれに合わせてバウンディングボックスの座標も変換されます。
(Albumentations公式サイト：https://albumentations.ai)

なお、画像データをバウンディングボックス付きで表示させるための関数もAlbumentationsの公式サイトに記載されていますので、こちらを活用させて頂きました。（ただし、オリジナルのコードはcoco形式に対応したものだったため、Pascal VOC形式に対応するように一部コードを修正しました）
(参考：https://albumentations.ai/docs/examples/example_bboxes/)

ソースコード

上記の【作業1】はGUIツールを用いてバウンディングボックスを手作業で付与していくため、ソースコードはありません。ここでは、【作業2】と【作業3】に相当するソースコードを記載します。

【作業2-1】アノテーションで作成したXMLファイルを取り込むためのクラス定義

import xml.etree.ElementTree as ET 

class xml2list(object):
    def __init__(self, classes):
        self.classes = classes

    def __call__(self, xml_path):
        ret = []
        xml = ET.parse(xml_path).getroot()

        for i in xml.iter("filename"):
            filename = i.text

        for size in xml.iter("size"):
            width = float(size.find("width").text)
            height = float(size.find("height").text)

        for obj in xml.iter("object"):
            difficult = int(obj.find("difficult").text)
            if difficult == 1:
                continue

            bndbox = [filename, width, height]
            name = obj.find("name").text.lower().strip() 
            bbox = obj.find("bndbox") 
            pts = ["xmin", "ymin", "xmax", "ymax"]

            for pt in pts:
                cur_pixel =  float(bbox.find(pt).text)
                bndbox.append(cur_pixel)

            label_idx = self.classes.index(name)
            bndbox.append(label_idx)
            ret += [bndbox]
        return np.array(ret) # [filename, width, height, xmin, ymin, xamx, ymax, label_idx]

上記のコードは、一つの画像データに犬や猫等のオブジェクトのうち、いずれか一つだけ写っている場合のみに対応しています。複数のオブジェクトのバウンディングボックス情報が１つのXMLファイル内に存在する場合は、コードの改良が別途必要になります。

【作業2-2】XMLファイル取り込み

import glob 
xml_paths = glob.glob("(XMLファイルが格納されているフォルダのパス)/*.xml")

category_names = ["dog", "cat"]  # カテゴリ名のリスト（XMLファイル内の<name>タグに記載されているカテゴリ名）

# XML取り込み用のクラスをインスタンス化
transform_anno = xml2list(category_names)

# 試しに先頭のXMLファイルの内容を読み込んでみる（下記コメントのような出力が得られればOK）
transform_anno(xml_paths[0])
# array([['cat.1.jpg', '300.0', '280.0', '32.0', '23.0', '300.0', '280.0', '1']], dtype='<U9')

【作業2-3】XMLファイルから取り込んだデータをDataFrame化

df = pd.DataFrame(columns=["filename", "width", "height", "xmin", "ymin", "xmax", "ymax", "category_id"])

for path in xml_paths:
    bboxs = transform_anno(path)

    for bbox in bboxs:
        tmp = pd.Series(bbox, index=["filename", "width", "height", "xmin", "ymin", "xmax", "ymax", "category_id"])
        df = df.append(tmp, ignore_index=True)

df = df.set_index("filename")
df = df.sort_values(by="filename", ascending=True)

display(df)

上記のコードが問題なく実行できれば、各画像ファイルに対するバウンディングボックスの情報を含んだ下記のようなデータフレームが表示されるはずです。

上記のデータフレームの「category_id」列の数値ですが、これはcategory_names変数（=["dog", "cat"] ）のインデックス番号をカテゴリのIDとして割り当てるようxml2listクラス内で処理しています。

【作業3-1】画像データをバウンディングボックス付きで表示する関数の定義

import random
import cv2
from matplotlib import pyplot as plt

BOX_COLOR = (255, 0, 0) # Red
TEXT_COLOR = (255, 255, 255) # White

def visualize_bbox(img, bbox, class_name, color=BOX_COLOR, thickness=2):
    """Visualizes a single bounding box on the image"""
    x_min, y_min, x_max, y_max = bbox  # バウンディングボックス形式：pascal_voc
    x_min, y_min, x_max, y_max = int(x_min), int(y_min), int(x_max), int(y_max)  # バウンディングボックス形式：pascal_voc

    # 左上の座標（x_min, y_min）から右下の座標(x_max, y_max)までの長方形を描画（＝バウンディングボックス）
    cv2.rectangle(img, (x_min, y_min), (x_max, y_max), color=color, thickness=thickness)

    ((text_width, text_height), _) = cv2.getTextSize(class_name, cv2.FONT_HERSHEY_SIMPLEX, 0.35, 1)    
    cv2.rectangle(img, (x_min, y_min - int(1.3 * text_height)), (x_min + text_width, y_min), BOX_COLOR, -1)

    # 画像データ中のバウンディングボックスにクラス名のテキストを追加
    cv2.putText(
        img,
        text=class_name,
        org=(x_min, y_min - int(0.3 * text_height)),
        fontFace=cv2.FONT_HERSHEY_SIMPLEX,
        fontScale=0.35, 
        color=TEXT_COLOR, 
        lineType=cv2.LINE_AA,
    )
    return img

def visualize(image, bboxes, category_ids, category_id_to_name):
    img = image.copy()
    for bbox, category_id in zip(bboxes, category_ids):
        class_name = category_id_to_name[category_id]
        img = visualize_bbox(img, bbox, class_name)
    plt.figure(figsize=(12, 12))
    plt.axis('off')
    plt.imshow(img)

【作業3-2】データ拡張実行

import os
import albumentations as A  # 使用したalbumentationsのバージョン:1.0.3
                            # （上記以外のバージョンだと下記のA.Compose辺りでエラーになるかも）

# データ拡張 定義
transform = A.Compose([
        A.HorizontalFlip(p=1),  # 水平方向移動
        A.ShiftScaleRotate(p=1),    # 回転
        A.RandomBrightnessContrast(p=1),    # コントラスト変更
        A.RGBShift(r_shift_limit=30, g_shift_limit=30, b_shift_limit=30, p=1),    #RGB変更
    ],
    bbox_params=A.BboxParams(format='pascal_voc', label_fields=['category_ids']),
)

jpg_paths = "（jpgファイルの入ったフォルダのパス）"
flist = os.listdir(jpg_paths)  # jpgファイルの一覧のリスト
image_list = []  # jpgファイルをndarray化したものを詰め込むためのリスト
category_id_to_name = {0: 'dog', 1: 'cat'}    # カテゴリのIDと名称の対応表


# フォルダ内のjpgファイルを順番にデータ拡張し、バウンディングボックス付きで可視化
for jpg_file in flist:
    print(jpg_file)
    image = cv2.imread(jpg_paths + "/" + jpg_file)
    image_list.append(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    # バウンディングボックスの座標データを取得
    bboxes_temp = list(df.loc[jpg_file, ["xmin", "ymin", "xmax", "ymax"]])
    print(bboxes_temp)
    bboxes = [tuple([float(bbox) for bbox in bboxes_temp])]
    # カテゴリIDを取得
    category_ids = [int(df.loc[jpg_file, "category_id"])]
    # print(image_list[-1])
    print(bboxes)
    print(category_ids)
    # データ拡張実行
    transformed = transform(image=image_list[-1], bboxes=bboxes, category_ids=category_ids)

    # データ拡張後の画像を表示
    visualize(
        transformed['image'],
        transformed['bboxes'],
        transformed['category_ids'],
        category_id_to_name
    )

上記処理はデータ拡張した画像を可視化するところまでで終わっており、データ拡張後の画像データの保存は行っていません。

さいごに

以上、物体検出のための前準備として、labelImgによるアノテーションとAlbumentationsによるデータ拡張作業の一連の流れについて説明しました。
少しでもご参考にして頂けましたら幸いです。