【超簡単】MMTrackingとYOLOv5を使って物体検出モデルの学習を教師データ作成含めてサクッと行う

公開日：2022年9月19日
最終更新日：2022年12月31日

はじめに
方法の概略
実行
環境構築
使用させて頂いたツール
その他

はじめに

以前からどうにか物体検出モデルの学習を簡単にできないかと模索してきました。

【pysot】【yolov5】2クラスの物体検出モデルの学習を教師データ作成含めてサクッと行うよ
 【pysot】【yolov7】2クラスの物体検出モデルの学習を教師データ作成含めてサクッと行うよ

今回以前以上に簡略化することに成功しました。

方法の概略

OpenMMLabの「MMTracking」と、pipでインストールできるYOLOv5を使わせて頂きました。

学習データとして準備するのは複数枚の画像？いや違います。
必要なのは目的の物体を撮影した動画ファイルです。

今回は各物体に対して20秒程度の動画を撮っただけです。あとでサンプルをダウンロードするのでどういった動画を撮ったかはそこで確認して下さい。

根気よく目的の物体を囲んでいく作業が必要？いや必要ありません。
準備した動画の数だけで良いです。サンプル動画（3クラス、4動画）を使うなら4回のみです。

実行

環境構築はこの記事の後半に載せておきます。

Clone

まずはリポジトリをクローンします。

git clone https://github.com/dai-ichiro/detection.git
cd detection

Gitが使えなければZIPでダウンロードして解凍しても問題ありません。
GitHubはこちらになります。

サンプル動画のダウンロード

python tools/download_sample_videos.py

これだけで「videos」フォルダが作成され中に動画がダウンロードされます。
どういった動画を撮ったかはここで確認して下さい。

3クラス、4動画になっています。
動画名がそのままクラス名になります。
1つのクラスに複数の動画を撮った場合にはフォルダにまとめます。フォルダ名がクラス名になります。

MOV動画になっていますがmp4でも問題ありません。

学習データ作成、学習

python detection.py --epochs 10

この1行で学習データ作成と学習が順に行われます。

ビデオの数だけ物体を囲む作業が必要です。
マウスで囲ってスペースキーを押して下さい。

引数は最低限のみ用意しました。

optional arguments:
  -h, --help            show this help message and exit
  --videos_dir VIDEOS_DIR
                        video folder name
  --epochs EPOCHS       total training epochs
  --batch BATCH         total batch size
  --weights WEIGHTS     initial weights path

videos_dirは動画ファイルの保存先です。デフォルトは「videos」になっています。
batchのデフォルトは8にしています。VRAMが不足するようなら減らす必要があります。
epochsのデフォルトは4と低めにしています。上記では10に設定しています。
weightsのデフォルトは「yolov5s.pt」にしています。

RTX 3080の環境では15～20分程度ですべての過程が終了しました。

VRAM 12GBのRTX 3080ではbatchはデフォルトのままで問題ありませんでした。
VRAM 6GBのRTX 3060 Laptopではbatchを2まで下げないと動作しませんでした。

python detection.py --epochs 10 --batch 2

推論

python inference.py --image sample_images/sample01.jpg --weights runs/train/exp/weights/best.pt

画像と学習済みweightsを指定して実行します。

学習済みweightsの保存先は異なる可能性があります。「exp」がすでに存在していれば「exp2」、「exp3」などになります。

これで記事の最初の写真が表示されました。なかなか精度は良さそうです。

（注意）「yolov5s.pt」というファイルがダウンロードされていますがそれは転移学習に使用したweightsです。今回の学習済みのweightsではないので注意して下さい。

環境構築

Windows 11
CUDA 11.6.2

Python 3.9.13

確実な動作を担保するためにそれぞれのバージョンを指定しています。
必ずしもこのバージョンでないと動かないというわけではありません。

pip install torch==1.12.1 torchvision==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu116
pip install mmcv-full==1.6.1 -f https://download.openmmlab.com/mmcv/dist/cu116/torch1.12.0/index.html
pip install mmdet==2.25.2
pip install mmtrack==0.14.0
pip install openmim==0.3.1
pip install mmengine==0.1.0
pip install yolov5==6.2.0
pip install gdown==4.5.1

「requirements.txt」を公開しているのでそちらを使うと以下のように簡単になります。

pip install torch==1.12.1 torchvision==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu116
pip install mmcv-full==1.6.1 -f https://download.openmmlab.com/mmcv/dist/cu116/torch1.12.0/index.html
pip install -r https://raw.githubusercontent.com/dai-ichiro/detection/main/requirements.txt

Python3.10で上記通りに環境構築しようとすると以下の警告が出てうまくいきません。

UserWarning: Failed to initialize NumPy: module compiled against API version 0x10 but this version of numpy is 0xf (Triggered internally at  ..\torch\csrc\utils\tensor_numpy.cpp:68.)