はじめに

YOLOで回転を含めたバウンディングボックス（Oriented Bounding Boxes, OBB)をTensorflow.jsで実行する手順のメモです。
この記事ではOBBの転移学習したモデルを出力するところまで。

間違っている可能性がありますので、自己責任で

手順

1. ラベリングデータセットを作成する
2. モデルの転移学習を実行する

ラベリングデータセットを作成する

画像のラベリングにはLabel Studioを使います。
Windowsのdockerを使った際の手順です。

Label Studioのセットアップ

dockerで準備するのが楽なので以下のコマンドで設定します。

docker pull heartexlabs/label-studio:develop

dockerイメージからコンテナを作成。
データを保存したいパスを設定。

docker run -it --name label-studio -p 8080:8080 -v <windows側のパス>:/label-studio/data heartexlabs/label-studio:develop

できたら localhost:8080 にアクセスする。

LabelStudioでプロジェクトの初期設定

最初のアクセスでローカルにアカウントを作る必要があるので
Sign in からEmailとPasswordを入力して作成。

プロジェクトの作成とラベリング

今回はバウンティングボックスに回転を追加したOBB(Oriented Bounding Boxes)のモデルを作成します。
OBBのラベルもObject Detaction with Bounding Boxesで作成できます。

Createでプロジェクトを作成します。
プロジェクトネームの設定やLabeling Setupラベルの種類を設定します。

画像をインポートしてラベリングを行っていきます。
回転を追加するときは以下の画像の赤丸をクリックしながら動かすとボックスが回転します。

ラベルの出力

ExportボタンからYOLOv8 OBBを選択して出力します。

モデルの転移学習を実行する

YOLOでモデルの転移学習を行います。
今回は環境をdockerで作成。

YOLOの環境構築

docker pull ultralytics/ultralytics:latest

jupyter labを使いたいので8888のポートを開けています。

docker run -it --ipc=host -p 8888:8888 --gpus all --name yolo ultralytics/ultralytics:latest 

(option)コンソールでjupyter labをインストールして起動

pip install jupyterlab
jupyter lab  --allow-root --ip 0.0.0.0

(option) データをtrain, val, testに分割する

YOLOで転移学習するにはデータをtrain, val, testに分割する必要がある。
Label Studioでは１つで出力されるのでjupyter labでデータを分割する。ChatGPTに聞いたコード

import os
import shutil
from sklearn.model_selection import train_test_split

# 画像とラベルのパスを取得
image_dir = 'stationary/images'
label_dir = 'stationary/labels'

images = [f for f in os.listdir(image_dir) if f.endswith(('.png', '.jpg', '.jpeg', '.webp'))]
labels = [f for f in os.listdir(label_dir) if f.endswith('.txt')]  # ラベルがテキストファイルだと仮定

# 両者が同じ数であり、ファイル名が対応していると仮定（例: image_1.jpg -> image_1.txt）
images.sort()
labels.sort()

# データを train, val, test に分割 (70% train, 15% val, 15% test の例)
train_images, test_images, train_labels, test_labels = train_test_split(
    images, labels, test_size=0.3, random_state=42)

val_images, test_images, val_labels, test_labels = train_test_split(
    test_images, test_labels, test_size=0.5, random_state=42)

# 分割したデータを保存するディレクトリを作成
train_dir = 'dataset/train'
val_dir = 'dataset/val'
test_dir = 'dataset/test'

os.makedirs(f'{train_dir}/images', exist_ok=True)
os.makedirs(f'{train_dir}/labels', exist_ok=True)
os.makedirs(f'{val_dir}/images', exist_ok=True)
os.makedirs(f'{val_dir}/labels', exist_ok=True)
os.makedirs(f'{test_dir}/images', exist_ok=True)
os.makedirs(f'{test_dir}/labels', exist_ok=True)

# 画像とラベルをそれぞれのディレクトリにコピー
def move_files(file_list, source_dir, target_dir):
    for file in file_list:
        shutil.copy(os.path.join(source_dir, file), os.path.join(target_dir, file))

move_files(train_images, image_dir, f'{train_dir}/images')
move_files(train_labels, label_dir, f'{train_dir}/labels')

move_files(val_images, image_dir, f'{val_dir}/images')
move_files(val_labels, label_dir, f'{val_dir}/labels')

move_files(test_images, image_dir, f'{test_dir}/images')
move_files(test_labels, label_dir, f'{test_dir}/labels')

print("データが train, val, test に分割されました！")

.yamlファイルを作成する

転移学習の設定などをまとめた.yamlを作成する。

path: ./dataset/
train: train
val: val
test: test

names:
    0: person
    1: car

転移学習を実行する

from ultralytics import YOLO
model = YOLO("yolo11n-obb.pt")
model.train(data="transfer.yaml", epochs=100, imgsz=640)

100回のエポックが終わると/ultralytics/runs/obb/train2/weights/best.pt が出力されているのでこのファイルをTensorflow.jsで使える形式に変換する。

custom_model = YOLO("/ultralytics/runs/obb/train2/weights/best.pt")
custom_model.export(format="tfjs")

これで出力される/ultralytics/runs/obb/train2/weights/best_web_model/内のデータをすべてコピーする。
このデータでTensorflow.jsで推論を行う。