YOLO用アノテーション確認ツール


YOLOの学習に使用するアノテーションファイルを参照し、タグ付けされた領域を切り出して表示するツールを作成しました。
きっと既に作られたものがあると思いましたが、探すよりは早いだろうと思い作ってしまいました。

アノテーションファイルの形式についてはこちらに解説されています。
作成したプログラムは以下のとおりです。
フォルダの中から再帰的に.txtファイルを検索し、対応する.jpgファイルがあれば、指定したクラスのタグが付けられた領域を切り出して表示するだけです。画像サイズは様々なので、MAX_SIZEに収まるようリサイズしています。

import os
import sys
import glob
import cv2

MAX_SIZE = 500
WIN_X = 300
WIN_Y = 200

def main(cls):
    txts = glob.glob("./**/*.txt", recursive=True)
    for txt in txts:
        jpg = txt.replace('.txt', '.jpg')
        if not os.path.exists(jpg):
            continue
        f = open(txt, 'r')
        lines = f.readlines()
        for line in lines:
            l = line.split()
            if l[0] == cls:
                print(jpg, line)
                img = cv2.imread(jpg)
                # タグ付けされた領域を切り出す
                h, w = img.shape[:2]
                cx = int(w * float(l[1]))
                cy = int(h * float(l[2]))
                rw = int(w * float(l[3]))
                rh = int(h * float(l[4]))
                rect = img[cy-int(rh/2):cy+int(rh/2), cx-int(rw/2):cx+int(rw/2)]
                # アスペクト比を維持して最大サイズ内に収める
                aspect = rw/rh
                if aspect > 1:
                    dst = cv2.resize(rect, dsize=(MAX_SIZE, int(MAX_SIZE/aspect)))
                else:
                    dst = cv2.resize(rect, dsize=(int(MAX_SIZE*aspect), MAX_SIZE))
                # 切り出した画像を表示
                cv2.imshow(jpg, dst)
                cv2.moveWindow(jpg, WIN_X, WIN_Y)
                cv2.waitKey(0)
                cv2.destroyAllWindows()
        f.close()

if __name__ == '__main__':
    args = sys.argv
    if 2 == len(args):
        main(args[1])
    else:
        print('Usage:',args[0],'class')

実行すると、以下のようにタグ付けされた領域(この例では「石祠」)が表示されます。