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

公開日：2022年7月16日
最終更新日：2022年9月20日

はじめに
詳細な手順
結果
拡張性
動作環境
Python環境構築
その他
2022年7月26日追記
2022年9月20日追記

はじめに

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

このような二つのパネルを検出します。
ペットボトルは大きさがわかりやすいように置いているだけです。

以前やったことと同じです。
touch-sp.hatenablog.com

Linuxコマンドをコピペすれば結果にたどり着くように工夫しています。

以前1クラスで同様のことをしています。今回は2クラスです。
touch-sp.hatenablog.com

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

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

詳細な手順

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」リポジトリ内に移動させます。

mv tracking_detection/make_dataset.py pysot/make_dataset.py

二つの動画ファイルをダウンロードして名前を変更して「pysot」フォルダに移動させます。

wget https://github.com/dai-ichiro/robo-one/raw/main/video_1.mp4
mv video_1.mp4 pysot/target.mp4
wget https://github.com/dai-ichiro/robo-one/raw/main/video_2.mp4
mv video_2.mp4 pysot/non_target.mp4

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

cd pysot
python make_dataset.py target.mp4 non_target.mp4

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

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

作成した「train.yaml」ファイルと「train_data」フォルダを「yolov5」リポジトリに移動させます。

mv train.yaml ../yolov5/train.yaml
mv train_data ../yolov5/train_data

「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

テスト画像をダウンロードします。

 wget https://github.com/dai-ichiro/robo-one/raw/main/test.jpg

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

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

結果

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

非常にうまくいっています。
動画ファイル名がクラス名になっています。

拡張性

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

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

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

動作環境

Ubuntu 22.04 on WSL2 (Windows11)
CUDA toolkit 11.3.1

python 3.10.4

Python環境構築

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

pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 --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.1.0
asttokens==2.0.5
backcall==0.2.0
beautifulsoup4==4.11.1
cachetools==5.2.0
certifi==2022.6.15
charset-normalizer==2.1.0
cycler==0.11.0
decorator==5.1.1
executing==0.8.3
filelock==3.7.1
fonttools==4.34.4
gdown==4.5.1
google-auth==2.9.1
google-auth-oauthlib==0.4.6
grpcio==1.47.0
idna==3.3
ipython==8.4.0
jedi==0.18.1
kiwisolver==1.4.4
Markdown==3.4.1
matplotlib==3.5.2
matplotlib-inline==0.1.3
numpy==1.23.1
oauthlib==3.2.0
opencv-python==4.6.0.66
packaging==21.3
pandas==1.4.3
parso==0.8.3
pexpect==4.8.0
pickleshare==0.7.5
Pillow==9.2.0
prompt-toolkit==3.0.30
protobuf==3.19.4
psutil==5.9.1
ptyprocess==0.7.0
pure-eval==0.2.2
pyasn1==0.4.8
pyasn1-modules==0.2.8
Pygments==2.12.0
pyparsing==3.0.9
PySocks==1.7.1
python-dateutil==2.8.2
pytz==2022.1
PyYAML==6.0
requests==2.28.1
requests-oauthlib==1.3.1
rsa==4.8
scipy==1.8.1
seaborn==0.11.2
six==1.16.0
soupsieve==2.3.2.post1
stack-data==0.3.0
tensorboard==2.9.1
tensorboard-data-server==0.6.1
tensorboard-plugin-wit==1.8.1
thop==0.1.1.post2207130030
torch==1.11.0+cu113
torchvision==0.12.0+cu113
tqdm==4.64.0
traitlets==5.3.0
typing_extensions==4.3.0
urllib3==1.26.10
wcwidth==0.2.5
Werkzeug==2.1.2
yacs==0.1.8