Ultralytics YOLOを学習する｜YAML設定から推論まで

YOLO

前回は、切り出した牌画像と4点アノテーションしたベース画像を使って、YOLOの学習データを自動生成するプログラムを作りました。

ここまでで、

画像を撮影
↓
矩形アノテーション
↓
牌画像を切り出す
↓
ベース画像を4点アノテーション
↓
YOLO学習データを自動生成

まで準備できました。

今回は、いよいよ作成したデータセットを使ってUltralytics YOLOを実際に学習します。

今回作成したtrain.pyでは、

YAMLファイルを読み込む
↓
学習パラメータを取得
↓
ベースモデルを読み込む
↓
YOLOを学習
↓
学習結果を保存
↓
test画像で推論
↓
推論画像・ラベル・confidenceを保存

までをまとめて実行しています。

YOLOの学習自体は、Ultralyticsを使うとかなりシンプルなコードで実行できます。

そこで今回は、学習処理だけを書くのではなく、データセットの場所や学習条件をYAMLファイルで管理し、学習からtestデータでの推論までまとめて実行できる形にしました。

この記事の後半には、今回使用しているtrain.pyとwork_tile.yamlの内容も掲載しています。

今回のプログラムでやること

今回のプログラムは、このYOLOモデル作成シリーズの最後となる学習部分です。

大きな流れは次のようになります。

work_tile.yaml
↓
データセットの場所・学習設定を読み込む
↓
ベースモデルを読み込む
↓
train / valデータで学習
↓
学習結果を保存
↓
test画像で推論
↓
推論結果を保存

前回作成したデータセットは、

images/
├── train/
├── val/
└── test/

labels/
├── train/
├── val/
└── test/

というYOLO用のフォルダ構成になっています。

今回は、このデータセットの場所や学習条件をYAMLファイルから指定します。

YAMLファイルでデータセットを設定する

今回使用しているwork_tile.yamlでは、データセットの場所を次のように指定しています。

path: ./02_augmentation/20260606

train: images/train
val: images/val
test: images/test

names:
  0: tile

pathを基準として、

./02_augmentation/20260606/images/train
./02_augmentation/20260606/images/val
./02_augmentation/20260606/images/test

の画像を使用します。

今回のYOLOでは、牌の種類まで分類するのではなく、

「画像内のどこに麻雀牌があるか」

を検出することを目的としています。

そのためクラスは、

names:
  0: tile

の1種類だけにしています。

牌の種類については、この後の別の画像分類モデルで判定する構成にしています。

学習設定もYAMLで管理する

今回は、YOLO標準のデータセット設定だけでなく、自分で使用する学習パラメータも同じYAMLファイルへまとめています。

train_param:
  epochs: 10
  imgsz: 640
  batch: 4
  device: mps
  base_model: "./basemodels/yolo11n.pt"
  result_root: "./train_data"

主な項目は次のとおりです。

項目内容
epochs学習を何epoch行うか
imgsz学習時の画像サイズ
batch一度に処理する画像枚数
device学習に使用するデバイス
base_modelベースとして使用するYOLOモデル
result_root学習結果の保存先

現在は、

epochs = 10
imgsz = 640
batch = 4

で設定しています。

また、Apple Silicon Macで実行しているため、

device: mps

としています。

学習条件をPythonコードの中に直接書かずYAMLへまとめておくことで、条件を変更するときにtrain.pyを書き換える必要がありません。

PythonからYAMLを読み込む

train.pyでは、最初にPyYAMLを使って設定ファイルを読み込みます。

yaml_path = "./workyaml/work_tile.yaml"

with open(yaml_path, "r") as f:
    cfg = yaml.safe_load(f)

その後、必要な学習設定を取り出します。

epochs = cfg["train_param"]["epochs"]
imgsz = cfg["train_param"]["imgsz"]
batch = cfg["train_param"]["batch"]
device = cfg["train_param"]["device"]
base_model = cfg["train_param"]["base_model"]

result_root = os.path.abspath(
    cfg["train_param"]["result_root"]
)

このように、

Python
→ 学習処理を書く

YAML
→ データセットや学習条件を書く

と役割を分けています。

今後、

  • epoch数を変更する
  • batchサイズを変更する
  • 使用するモデルを変更する
  • 学習データを変更する

といった場合も、基本的にはYAML側を変更するだけで対応できます。

ベースモデルを読み込む

モデルの読み込みには、UltralyticsのYOLOクラスを使用します。

model = YOLO(base_model)

今回の設定では、

base_model: "./basemodels/yolo11n.pt"

としているため、yolo11n.ptをベースモデルとして使用しています。

YOLOを学習する

実際に学習する部分はかなりシンプルです。

model.train(
    data=yaml_path,
    epochs=epochs,
    imgsz=imgsz,
    batch=batch,
    device=device,
    project=result_root,
    name=WORK,
)

ここで、

data
epochs
imgsz
batch
device

などを指定しています。

dataには先ほどのwork_tile.yamlを渡しているため、

train
val
names

などのデータセット情報をUltralytics側が読み込んで学習してくれます。

ここまで自分で学習データを作る部分はかなり長かったですが、実際にYOLOを学習するコード自体は非常に短いです。

学習結果を日時ごとに保存する

学習条件を変えながら何度も実行すると、

「どのフォルダがどの学習結果なのか」

分からなくなりやすくなります。

そこで今回は、実行した日時をフォルダ名へ入れています。

timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

WORK = f"work_{timestamp}"

例えば、

work_20260830_201530

のような名前になります。

学習結果は、

train_data/
└── work_20260830_201530/

のように、実行ごとに別のフォルダへ保存されます。

学習条件を変更しながら結果を比較するときにも便利です。

学習後にtest画像で推論する

学習が完了したら、そのままtestフォルダの画像を使って推論します。

まず、YAMLからtest画像の場所を取得します。

path = cfg["path"]
test = cfg["test"]

test_dir = f"{path}/{test}"

その後、model.predict()を実行します。

model.predict(
    source=test_dir,
    imgsz=640,

    save=True,
    save_txt=True,
    save_conf=True,

    project=f"{result_root}/{WORK}",
    name="predict",

    exist_ok=True
)

今回は、

save=True

で検出結果を描画した画像を保存し、

save_txt=True

でYOLO形式の検出結果を保存しています。

さらに、

save_conf=True

としているため、confidenceもTXTファイルへ保存されます。

これによって、学習が終了したあとに別のプログラムを実行しなくても、そのままtestデータで検出結果を確認できるようにしています。

以下はtestデータを推論した画像です。青い枠線で囲われたところが牌として検出した結果です。

学習から推論までまとめて実行する理由

YOLOを学習するだけなら、もっと短いコードでも実行できます。

ただ、今後は、

  • データセットを変更する
  • epochsを変更する
  • batchを変更する
  • ベースモデルを変更する
  • 学習結果を比較する

といった作業を何度も行うことになります。

そこで今回は、

YAMLを変更
↓
train.pyを実行
↓
学習
↓
結果を日時付きフォルダへ保存
↓
testデータで推論

までをまとめました。

これによって、条件を変えて再学習するときも、基本的にはYAMLを書き換えてプログラムを実行するだけで済みます。

完成したコード

今回使用したtrain.pyの全文です。

コードが長いため、デフォルトでは折りたたんでいます。必要に応じて「train.py の全文を見る」を「work_tile.yamlの全文を見る」クリックしてください。

train.py の全文を見る 
import yaml
import os
from ultralytics import YOLO
from datetime import datetime


yaml_path = "./workyaml/work_tile.yaml"

with open(yaml_path, "r") as f:
    cfg = yaml.safe_load(f)


path = cfg["path"]
test = cfg["test"]
test_dir = f"{path}/{test}"

epochs = cfg["train_param"]["epochs"]
imgsz = cfg["train_param"]["imgsz"]
batch = cfg["train_param"]["batch"]
device = cfg["train_param"]["device"]
base_model = cfg["train_param"]["base_model"]

result_root = os.path.abspath(
    cfg["train_param"]["result_root"]
)

print(result_root)

# モデル読み込み
model = YOLO(base_model)

timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

WORK = f"work_{timestamp}"

model.train(
    data=yaml_path,
    epochs=epochs,
    imgsz=imgsz,
    batch=batch,
    device=device,
    project=result_root,
    name=WORK,
)


model.predict(
    source=test_dir,
    imgsz=640,

    save=True,
    save_txt=True,
    save_conf=True,

    project=f"{result_root}/{WORK}",
    name="predict",

    exist_ok=True
)
work_tile.yaml の全文を見る 
path: ./02_augmentation/20260606

train: images/train
val: images/val
test: images/test

names:
  0: tile


train_param:
  epochs: 10
  imgsz: 640
  batch: 4
  device: mps
  base_model: "./basemodels/yolo11n.pt"
  result_root: "./train_data"

実行方法

最後に、今回のYOLO学習プログラムの実行方法を紹介します。

必要なライブラリ

今回使用する主なライブラリは、

  • Ultralytics
  • PyYAML

です。

インストールされていない場合は、

pip install ultralytics pyyaml

を実行します。

学習データを用意する

前回作成したYOLO用データセットを用意します。

02_augmentation/
└── 20260606/
    ├── images/
    │   ├── train/
    │   ├── val/
    │   └── test/
    │
    └── labels/
        ├── train/
        ├── val/
        └── test/

画像とラベルのファイル名は対応させておきます。

例えば、

images/train/base_001_000.jpg
labels/train/base_001_000.txt

という形です。

YAMLを設定する

次に、work_tile.yamlへデータセットの場所と学習条件を記述します。

例えば学習回数を変更したい場合は、

epochs: 100

のように変更します。

使用するベースモデルを変更する場合は、

base_model: "./basemodels/yolo11n.pt"

を書き換えます。

プログラムを実行する

設定が終わったら、ターミナルから、

python train.py

を実行します。

環境によっては、

python3 train.py

で実行します。

学習が開始され、学習終了後にはtest画像への推論も自動で実行されます。

結果は、

train_data/
└── work_YYYYMMDD_HHMMSS/
    ├── 学習結果
    └── predict/
        └── 推論結果

のように、実行日時ごとに保存されます。

まとめ

今回は、これまで作成してきたYOLO学習データを使って、Ultralytics YOLOを実際に学習するところまで進めました。

今回のプログラムでは、

YAMLから設定を読み込む
↓
YOLOモデルを読み込む
↓
train / valデータで学習
↓
日時ごとに学習結果を保存
↓
test画像で推論
↓
画像・ラベル・confidenceを保存

までをまとめて実行しています。

これで、今回のYOLOモデル作成シリーズで進めてきた、

画像を撮影
↓
矩形アノテーション
↓
牌画像を切り出す
↓
ベース画像を4点アノテーション
↓
学習データを自動生成
↓
YOLOを学習

という一連の流れが完成しました。

最初は、撮影した画像をそのままアノテーションして学習していました。

そこから、

「もっと効率よく学習データを増やせないか」

と考え、

牌画像を素材として切り出す
↓
ベース画像へ射影変換して貼り付ける
↓
YOLOラベルも自動生成する
↓
データ拡張する

という独自のデータセット生成方法へ変更しました。

最終的には、一度素材を準備しておけば、そこから大量の学習データを生成し、そのままYOLOの学習とtestデータでの推論まで進められるようになりました。

今後は、実際の認識結果を確認しながら、データセットや学習条件を調整して精度を改善していきます。

関連記事

コメント

タイトルとURLをコピーしました