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

公開日：2022年4月8日
最終更新日：2022年9月20日

はじめに
詳細な手順
結果
拡張性
動作環境
Python環境構築
その他
2022年4月9日追記：ノートPCで試す
2022年4月26日追記：Ubuntu 22.04で試す
2022年9月20日追記：最新記事はこちら

はじめに

タイトルにあるように教師データ作成を含めて物体検出モデルの学習をサクッと行います。

今回はこのようなアンパンマンの人形を検出してみます。

動画撮影してそれをObject Trackingすることによって教師データを自動で作成します。

携帯電話（iPhone 7）で画質を720p/30fps（1280×720）に設定して30秒ほどの動画を撮りました。

どのような動画を撮ったかは後でダウンロードするので実際に見てみて下さい。

詳細な手順

Pythonの環境構築はこのブログの後半を参照して下さい。

環境構築後は以下を順に実行するだけで教師データ作成→学習→推論ができます。

3つのリポジトリをクローンします。
3番目が今回のために作ったリポジトリになります。

git clone https://github.com/STVIR/pysot
git clone https://github.com/ultralytics/yolov5
git clone https://github.com/dai-ichiro/tracking_detection

「tracking_detection」リポジトリ内の「make_dataset.py」を「pysot」リポジトリ内にコピーします。

cp tracking_detection/make_dataset.py pysot/make_dataset.py

「tracking_detection」リポジトリ内の「download_sample_video.py」を実行して「sample.MOV」をダウンロードします。そしてそれを「pysot」リポジトリ内に移動させます。

python tracking_detection/download_sample_video.py
mv sample.MOV pysot/sample.MOV

「pysot」リポジトリ内に入って「sample.MOV」を引数として「make_dataset.py」を実行します。

cd pysot
python make_dataset.py sample.MOV

動画の1番目のフレームが表示されるのでターゲットになる物体を選択しエンターキーまたはスペースキーを押します。

その後Object Trackingが始まり「train.yaml」ファイルと「train_data」フォルダが作成されます。

作成した「train.yaml」ファイルと「train_data」フォルダを「yolov5」リポジトリにコピーします。

cp train.yaml ../yolov5/train.yaml
cp train_data ../yolov5/train_data -r

「yolo5」リポジトリに入って「train.py」を実行します。

cd ../yolov5
python train.py --batch 8 --epochs 10 --data train.yaml --weights yolov5s.pt

これによって「best.pt」が作成されます。
「best.pt」は「runs/train/exp/weights」フォルダ内に保存されています。
「exp」フォルダがすでに存在している場合はフォルダ名が「exp1」になったり「exp2」になったりします。

ここから先はどのディレクトリでも実行可能です。
「trackin_detection」内の「inference.py」と先ほど保存された「best.pt」が必要になります。
いったんすべてのリポジトリから抜けて「best.pt」を現在のディレクトリにコピーします。

cd ../
cp yolov5/runs/train/exp/weights/best.pt best.pt

「tracking_detection」リポジトリ内の「download_sample_picture.py」を実行して「test.jpg」をダウンロードします。

python tracking_detection/download_sample_picture.py

「tracking_detection」リポジトリ内の「inference.py」を実行します。その際に引数として「test.jpg」と「best.pt」を渡します。

python tracking_detection/inference.py test.jpg best.pt

結果

このような画像が表示されます。

30秒のビデオを撮影しただけで非常にうまくいきました。
動画ファイル名が「sample.MOV」だったのでクラス名が「sample」になっています。

拡張性

今回はクラス数 1、動画 1でした。

複数のクラスに対応しています。「make_dataset.py」の引数を増やすだけです。

1つのクラスに複数の動画を撮った場合にも対応しています。その際は同一クラスの動画を一つのフォルダにいれてフォルダ名を引数にして下さい。

動画ファイルを引数とした場合はそのファイル名がクラス名になります。
フォルダを引数とした場合にはフォルダ名がクラス名になります。

動作環境

Intel Core i7-11700
GTX 1080 (VRAM 8GB)

Ubuntu 20.04 on WSL2 (Windows11)
CUDA toolkit 11.3

WSL バージョン: 0.56.2.0
カーネル バージョン: 5.10.102.1
WSLg バージョン: 1.0.30
MSRDC バージョン: 1.2.2924
Direct3D バージョン: 1.601.0
Windows バージョン: 10.0.22000.593

python 3.9.5

Python環境構築

すべてpipでインストール可能です。

pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu113
pip install -r https://raw.githubusercontent.com/ultralytics/yolov5/master/requirements.txt
pip install yacs
pip install gdown

absl-py==1.0.0
beautifulsoup4==4.10.0
cachetools==5.0.0
certifi==2021.10.8
charset-normalizer==2.0.12
cycler==0.11.0
filelock==3.6.0
fonttools==4.31.2
gdown==4.4.0
google-auth==2.6.2
google-auth-oauthlib==0.4.6
grpcio==1.44.0
idna==3.3
importlib-metadata==4.11.3
kiwisolver==1.4.2
Markdown==3.3.6
matplotlib==3.5.1
numpy==1.22.3
oauthlib==3.2.0
opencv-python==4.5.5.64
packaging==21.3
pandas==1.4.2
Pillow==9.1.0
pkg_resources==0.0.0
protobuf==3.20.0
pyasn1==0.4.8
pyasn1-modules==0.2.8
pyparsing==3.0.7
PySocks==1.7.1
python-dateutil==2.8.2
pytz==2022.1
PyYAML==6.0
requests==2.27.1
requests-oauthlib==1.3.1
rsa==4.8
scipy==1.8.0
seaborn==0.11.2
six==1.16.0
soupsieve==2.3.1
tensorboard==2.8.0
tensorboard-data-server==0.6.1
tensorboard-plugin-wit==1.8.1
thop==0.0.31.post2005241907
torch==1.11.0+cu113
torchvision==0.12.0+cu113
tqdm==4.64.0
typing_extensions==4.1.1
urllib3==1.26.9
Werkzeug==2.1.1
yacs==0.1.8
zipp==3.8.0

その他

今回の内容は過去の記事の延長です。
touch-sp.hatenablog.com
touch-sp.hatenablog.com
Linuxコマンドのみですべてが完結するように工夫しました。
Pythonスクリプトを修正する必要はありません。

PythonスクリプトはGitHubで公開しています。
GitHub - dai-ichiro/tracking_detection