YOLOv7を使って自作データセットで物体検出してみた

2022.07.19

こんちには。

データアナリティクス事業本部機械学習チームの中村です。

YOLOv7の論文が2022-07-06にarXivに公開されたようですね🎉🎉🎉

ソースコードもGitHubで公開されています。

せっかくなので今回は、以下の記事で紹介した自作データのトレーニングを、YOLOv7を使ってやってみたいと思います。

YOLOv7の概要

YOLOv7は、YOLOv4、Scaled-YOLOv4, YOLORと同じグループが開発した最新の物体検出処理です。

MS COCOデータセットでは既存手法を大きく上回る結果と報告されています。

ざっと見たところポイントは以下となっています。

  • concatenateモデルのアーキテクチャを進化させたELANおよびE-ELANを使用
    • concatenateモデルはDenseNetやCSPVoVNetなどのようにaddの代わりにconcatするモデルのこと
  • concatenateモデルに適した複合スケーリング方法の提案
  • 近年の研究成果である以下を適切な形で設計して盛り込んだ
    • ラベル割り当て戦略の最適化(OTAなど)
    • re-parameterizationの適用(RepVggなど)

ラベル割り当て戦略の最適化は、正解データをどの場所のbounding boxに割り当てればよいかを最適化する手法です。

詳細は以下の論文等を参考にされてください。

re-parameterizationは、skip-connectionなどを持つ複数のconvとbatch normalizationを、一つのbias付きconvに置き換える手法です。

詳細は以下の論文等を参考にされてください。

YOLOv7の詳細については今後ブログで紹介できればと思います。

実行環境

今回はGoogle Colaboratory環境で実行しました。

ハードウェア情報は以下の通りです。

  • GPU: Tesla P100 (GPUメモリ16GB搭載)
  • メモリ: 26GB

主なソフトウェア・ライブラリのバージョンは以下となります。

  • CUDA: 11.2
  • PyTorch: 1.12.0+cu113

またデータを永続化するために、Google Drive上に作業ディレクトリを実施します。

この記事の通りに学習すると、容量をかなり使用するため空き容量を考慮して実行してください。(私がやったときは80GB程度空きがあれば大丈夫でした)

使用するモデル

YOLOv7は以下のリンクにあるように複数の規模のモデルがあります。

今回はこの中で、YOLOv7-E6を使用してみました。

データセットの準備

参考:データを準備することが難しい場合

以下にマスク着用を判定するデータセットがありますので、そちらを使用して試されてください。

今回使用したデータ

以下の記事で準備したデータセットを使用します。

具体的には、以下のようなメロンソーダアメとオレンジアメが写った写真となります。

アノテーション結果も同じものを使用します。

YOLOv7は読み込み対象がYOLOフォーマットとなっていますので、Export時にYOLOフォーマットを指定して出力します。

最終的には、以下のようなフォルダ構造となるように構成します。

└─candy
    ├─train
    │  ├─images
    │  │      2022-07-07_07-40-32.jpg
    │  │      ...
    │  └─labels
    │          2022-07-07_07-40-32.txt
    │          ...
    └─valid
        ├─images
        │      2022-07-07_09-43-43.jpg
        │      ...
        └─labels
                2022-07-07_09-43-43.txt
                ...

データセットの情報はこのようになります(前回と同じ)。

set JPGファイル数 物体数(melon-soda) 物体数(orange)
train 12 28 153
valid 6 4 29

Google Colaboratoryでの実行

セットアップ

データを永続化するため、ストレージはGoogle Driveを使用します。

実行前に、Google Driveをマウントします。

そして以下のコマンドで作業ディレクトリをGoogle Driveに変更します。

%cd /content/drive/MyDrive

各種ライブラリをインストールします。まずはGitHubからコードを取得します。

!git clone https://github.com/WongKinYiu/yolov7
%cd yolov7

動作させた時点のコミットは以下となります。

Commit: e3186245c37129d8c3a8c04a5de62cde8f9abf97
Parents: 987dd8324b0a28a6769f9ecafda31561fb55e5e4
Author: Kin-Yiu, Wong <102582011@cc.ncu.edu.tw>
Committer: GitHub <noreply@github.com>
Date: Mon Jul 18 2022 11:11:19 GMT+0900 (日本標準時)

yolov7内にrequirements.txtがありますが、Colabに元から入っているパッケージと衝突などが発生したため、PyYAMLのみインストールします。

具体的には、requirements.txtを以下のように編集します。

# # Usage: pip install -r requirements.txt

# # Base ----------------------------------------
# matplotlib>=3.2.2
# numpy>=1.18.5
# opencv-python>=4.1.1
# Pillow>=7.1.2
PyYAML>=5.3.1
# requests>=2.23.0
# scipy>=1.4.1
# torch>=1.7.0,!=1.12.0
# torchvision>=0.8.1,!=0.13.0
# tqdm>=4.41.0
# protobuf<4.21.3

# # Logging -------------------------------------
# tensorboard>=2.4.1
# # wandb

# # Plotting ------------------------------------
# pandas>=1.1.4
# seaborn>=0.11.0

# # Export --------------------------------------
# # coremltools>=4.1  # CoreML export
# # onnx>=1.9.0  # ONNX export
# # onnx-simplifier>=0.3.6  # ONNX simplifier
# # scikit-learn==0.19.2  # CoreML quantization
# # tensorflow>=2.4.1  # TFLite export
# # tensorflowjs>=3.9.0  # TF.js export
# # openvino-dev  # OpenVINO export

# # Extras --------------------------------------
# ipython  # interactive notebook
# psutil  # system utilization
# thop  # FLOPs computation
# # albumentations>=1.0.3
# # pycocotools>=2.0  # COCO mAP
# # roboflow

そしてパッケージをインストールをします。

!pip install -r requirements.txt

学習済みモデルの取得

YOLOv7の学習済みモデルを取得します。以下にダウンロード可能なモデル一覧があります。

本記事執筆時点では、以下のリンクでダウンロードできます。

(まだYOLOv7自体がpre-releaseなのでリンクは今後変わる可能性があります。)

!wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7-e6.pt -P ./checkpoints

一応、ダウンロードしたモデルで動作確認用に推論してみましょう。

!python detect.py --source inference/images/horses.jpg \
  --weights checkpoints/yolov7-e6.pt \
  --conf 0.25 --img-size 1280 --device 0
from IPython.display import Image
Image("runs/detect/exp2/horses.jpg")

以下のお馬さん画像が物体検出されました。

データセットの配置

アノテーション済みのデータセットを以下のフォルダに配置します。

dataset/candy

configファイルの修正

データ用のconfigファイルをdata/candy.yamlに作成して、以下を記述します。

特に難しい記述は不要な印象です。

# COCO 2017 dataset http://cocodataset.org

# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: ./dataset/candy/train
val: ./dataset/candy/valid

# number of classes
nc: 2

# class names
names: ['melon-soda', 'orange']

学習の実行

configファイルの修正ができれば、以下で学習を行うことができます。

YOLOv7-E6などhead部の最終段でIAuxDetectを使用するモデルは、train_aux.pyで動かす必要がありますのでご注意ください。

(逆にYOLOv7-Xなどは、IAuxDetectではなくIDetectとなっているため、train.pyで学習することが可能です。)

!python train_aux.py --workers 2 --batch-size 8 \
  --data data/candy.yaml \
  --cfg cfg/training/yolov7-e6.yaml \
  --weights 'checkpoints/yolov7-e6.pt' \
  --name yolov7-e6-candy \
  --hyp data/hyp.scratch.p6.yaml \
  --epochs 300 \
  --device 0

またコマンドの--workers 2 --batch-size 8については、高すぎるとGPUメモリ不足で動かなくなる可能性があります。必要に応じて調整してください。

前述の実行環境では約80分程度で学習が終わりました。

結果は以下のパスに配置されます。(前述の--nameオプションが出力先のフォルダ名になっています)

runs/train/yolv7-e6-candy

推論テスト

まずvalid用のデータで推論してみます。

!python detect.py --weights runs/train/yolov7-e6-candy/weights/best.pt \
  --conf 0.25 --img-size 640 \
  --source dataset/candy/valid/images \
  --name exp-yolov7-e6-candy-valid

結果は以下のフォルダに配置されます。

runs/detect/exp-yolov7-e6-candy-valid

以下が推論結果のサンプルです。

また動画でも同様のコマンドで推論できるようになっています。

dataset/candy/testに動画データをおいていただき、以下のコマンドを実行します。

!python detect.py --weights runs/train/yolov7-e6-candy/weights/best.pt \
  --conf 0.25 --img-size 640 \
  --source dataset/candy/test \
  --name exp-yolov7-e6-candy-test

動画は処理結果をお見せできることができませんが、きちんと動作することが確認できました。

まとめ

いかがでしたでしょうか。

かなり簡単にカスタムデータを学習できる構成となっていることが分かりました。

ここからチューニング等を実施していくためには論文やソースコードを理解していくことが必要そうです。

今後のブログ記事で、YOLOv7の論文に記載されている理論やソースコードの実装をついても取り扱っていけたらと思います。