はじめに
機械学習について勉強を始め、その一環で画像から物体検出を実装しようと思いyoloを触ってみました。
チェスの駒を題材にしたのは、画像からチェスの盤面を理解し、最終的にロボットアームと組み合わせて自動チェスロボットを自作したいと思っているからです。
環境
ubuntu24.04
Python 3.11.11
GPU: RTX3070
YOLOについて
詳しいことは多く記事が出ているので、今回は割愛します。
バージョンに関しては実装時点での最新バージョンYOLO11を使います。
（2024年にYOLOv9、YOLOv10、YOLO11の3バージョン出ていたみたいです）
この記事によると新しいバージョンのほうが認識精度、速度ともに上がっているみたいです。
環境構築
YOLOの必要環境として、Python3.8以上、Pytorch1.8以上が必要なので、それらをインストールします。
これは以前自分の記事で行ったので、良かったら参考にしてください。（宣伝）
YOLOのインストール
pip install ultralytics
参考：公式GitHub,公式クイックスタート
インストールされたらテストコードを動かしてみます。
yolo predict model=yolo11n.pt source='https://ultralytics.com/images/bus.jpg'
すると実行したディレクトリ直下にruns/detect/predict/ディレクトリが生成され、検出を行う画像bus.jpgがダウンロードされます。
runs/detect/predict/にあるbus.jpgが検出結果になります。
左が検出前の画像、右が検出後の画像
これでYOLOが正しく動作していることが分かります。
YOLOでの学習
学習にはこちらの記事を参考にさせていただきました。
学習用のモデルの調達
自前でデータセットを用意するのは大変なので、今回はインターネット上のデータセットを使用します。
データセットはこちらを使用しました。
Download Projectsを押し、今回はYOLO11を使うのでYOLOv11をクリック
今回はデータセットをローカルにダウンロードして学習するので、Download zip to computerをクリックしてContinueを押す。
そうすると学習用のデータセットをダウンロードすることができます。
（要アカウント登録）
学習
今回の学習は以下のコードで行いました。
dataは学習に使うデータセット、epochsはどのくらい学習するか（訓練データを何回見るか）です。
dataは適切なパスに、epochsも適切な回数に設定してください
（多すぎると過学習で、少ないと学習が足らずに良くないみたいです）
from ultralytics import YOLO
# ベースとするモデル
model = YOLO('yolo11n.pt')
# M1 macのGPUを使ってモデルを学習
results = model.train(
data='datasets/Chess Piece Detection.v2i.yolov11/data.yaml',
epochs=50,
imgsz=640,
device='0'
)
また、ダウンロードしたデータセットの中のdata.yamlは次のようになっています。
train: ../train/images
val: ../valid/images
test: ../test/images
nc: 12
names: ['black-bishop', 'black-king', 'black-knight', 'black-pawn', 'black-queen', 'black-rook', 'white-bishop', 'white-king', 'white-knight', 'white-pawn', 'white-queen', 'white-rook']
roboflow:
workspace: chess-piece-detection-lydqy
project: chess-piece-detection-5ipnt
version: 2
license: CC BY 4.0
url: https://universe.roboflow.com/chess-piece-detection-lydqy/chess-piece-detection-5ipnt/dataset/2
これの
train: ../train/images
val: ../valid/images
test: ../test/images
を適切なパスに設定してください。
ここで少しつまずいたので、エラーを見ながら設定するといいと思います。
これらの設定が終わったら、実行すると学習が始まり、以下のような進捗が見られます。
Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size
19/50 4.27G 1.139 0.5812 1 32 640:
Class Images Instances Box(P R mAP50
all 1592 30400 0.971 0.966 0.988 0.759
Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size
20/50 3.5G 1.131 0.576 0.9961 44 640:
Class Images Instances Box(P R mAP50
all 1592 30400 0.974 0.969 0.989 0.76
Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size
21/50 3.56G 1.125 0.5672 0.9952 32 640:
Class Images Instances Box(P R mAP50
all 1592 30400 0.976 0.97 0.989 0.758
Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size
22/50 3.5G 1.114 0.5595 0.9915 19 640:
Class Images Instances Box(P R mAP50
all 1592 30400 0.973 0.969 0.989 0.762
Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size
23/50 3.98G 1.105 0.5528 0.9896 433 640:
推論
学習が終わるとruns/detext/trainディレクトリが生成されていると思います。そこのweight/best.ptが学習結果となります。
これを使って、以下のコードで推論できます。
from ultralytics import YOLO
# Load a model
model = YOLO("best.pt") # pretrained YOLO11n model
# Run batched inference on a list of images
results = model.predict("input_file_path.jpg") # return a list of Results objects
# Process results list
for result in results:
result.save(filename="result.jpg") # save to disk
推論の結果は以下になります。
左が元画像で右が推論結果
終わりに
YOLOはフレームワークが出来上がっているので、簡単に使えて便利だと感じました。
この記事が参考になれば幸いです。