UltralyticsでBBOXアノテーションファイルを作成、CVATでアノテーションインポート

アノテーションを半自動化したい！

煌びやかなAIの裏には、人間の地味な作業の積み重ねがある。
そんなことを思いながらアノテーションを作業をしていましたが、いい加減疲れました。
幸いこれまで学習してきたことで物体検出のカスタムモデルもそこそこの検出精度を持っています。
このカスタムモデルを使って、8割くらいはアノテーションが終わっている状態にして、人間様は微調整をするだけにしたい！
アノテーション作業に利用しているCVATには、モデルを読み込んでアノテーションを行う機能があります。
次の記事を参考にやろうと思ったのですが、後半のyamlやmain.pyあたりで挫折しました・・・

CVATではアノテーションのインポートもできるので、

• YOLOv8でアノテーションファイルを作成
• アノテーションファイルをCVATにインポート

という手順でアノテーションの半自動化を泥臭く実現しましたので、その手順をご紹介いたします。

Ultralyticsによるbboxなアノテーションファイルの作成

Ultralyticsにはたくさん便利ツールが準備されています。
coco形式からyolo形式に変えるツールとかも使わせていただきました。
なら、YOLOを利用したアノテーション作成もあるだろう！ → ありました！（auto_annotate）

あったんですが、セグメンテーションのツールでした。
私は物体検出だけできればいいので、BBOXで終わりでよかったんですが、気を利かせてセグメンテーションまでしてくれます。
このソースを参考につくり直せば、とも思いましたが、脳みそが停止していたので、Windows Copilotにお願いして作ってもらいました。

Windows Copilotへのプロンプト（参考）

UltralyticsのYOLOを利用して、既に学習させたモデルファイルで推論を行い、
指定フォルダにある画像のアノテーションファイルを自動で作成したいと思っています。 
アノテーションファイルは、YOLO形式とし、BBOXでの指示したいと思っています。
画像フォルダの中の対象ファイルは、png、jpgとしてサンプルを提示いただけないでしょうか

拡張子の指定方法がちょっと大丈夫かな？という感じはしましたが、少なくとも主な対象しているPNGファイルでは動作したので、ほぼそのまま利用しています。
指定したフォルダに、画像と同名のtxtファイルが生成されます。
既存のアノテーションと比較したかったので、YOLO形式としています。
predictの中は、実際にモデル利用時の、imgszや confを設定ください。

auto_annotation.py

import os
from pathlib import Path
from ultralytics import YOLO

# モデルのロード
model = YOLO('path_to_your_trained_model.pt')  # 学習済みモデルのパスを指定します

# 画像フォルダの指定
image_folder = 'path_to_your_image_folder'
output_folder = 'path_to_output_folder'  # 出力フォルダを指定します

# 指定した出力フォルダが存在しない場合、フォルダを作成
os.makedirs(output_folder, exist_ok=True)

# 画像フォルダ内のPNGおよびJPG画像に対して推論を実行
for image_path in Path(image_folder).glob('*.[pjJ][pnN][gG]'):
    # 画像の推論
    results = model.predict(str(image_path))

    # 各画像ごとにアノテーションファイルを作成
    for result in results:
        # バウンディングボックスの情報を取得
        bboxes = result.boxes.xyxy.cpu().numpy()
        classes = result.boxes.cls.cpu().numpy()
        
        file_name = os.path.splitext(os.path.basename(image_path))[0]
        
        # アノテーションファイルのパス
        annotation_file = os.path.join(output_folder, file_name + '.txt')
        
        # アノテーションファイルの作成
        with open(annotation_file, 'w') as f:
            for bbox, cls in zip(bboxes, classes):
                class_id = int(cls)
                x_center = (bbox[0] + bbox[2]) / 2 / result.orig_shape[1]
                y_center = (bbox[1] + bbox[3]) / 2 / result.orig_shape[0]
                width = (bbox[2] - bbox[0]) / result.orig_shape[1]
                height = (bbox[3] - bbox[1]) / result.orig_shape[0]
                # YOLO形式で保存 (class_id, x_center, y_center, width, height)
                f.write(f"{class_id} {x_center} {y_center} {width} {height}\n")

print("アノテーションファイルの作成が完了しました！")

CVATへアノテーションファイルをインポートする

CVATからアノテーションをエクスポートされる形と同じ構造にすることで、アノテーションをインポートすることができます。
CVATからエクスポートしたファイルを参考にして作成するのがいいと思います。
具体的には次のような形にします。

• 全体をzip圧縮する
• フォルダ構造は次のようにする。

/
└ labels
  └ train
    └ annotation_1.txt
      annotation_2.txt
      ・・・
data.yaml
train.txt

data.yaml、train.txtは次のように書きます。
私の物体検出の独自モデルには、クラスが１つしかないので、0:pc_nameだけとなっています。
train.txtは、ファイルがpngなど画像ファイルなのが落とし穴です。

data.yaml

names:
  0: pc_name
path: .
train: train.txt

train.txt

data/images/train/annotation_1.png
data/images/train/annotation_2.png
・・・

train.txtの画像のリストを作成する時には、CVATにインポートした画像ファイルのあるフォルダで、コマンドプロンプト dir /s /b > train.txtとすると、一覧がテキスト出力されます。あとはテキストエディタの置換で整形します。

使ってみての感想

うきょー！アノテーションらくちんー！