[DSW Gallery] Example of image detection based on YOLOX model

本記事では、YOLOX モデルを例に、EasyCV を使った物体検出モデルのトレーニングと予測の方法を紹介します。
  YOLO シリーズの開発は、主に V1 (2015)、V2 (2016)、V3 (2018)、V4 (2020)、V5 (2020) といったバージョンが含まれます。YOLO は、2015 年に Joseph Redmon と Ross Girshick が You Only Look Once: Unified, Real-Time Object Detection (Link) で提案した one-stage 物体検出ネットワークです。YOLO シリーズは、高速な応答、高い精度、シンプルな構造、デプロイのしやすさから、エンジニアや研究者に広く利用されています。一方で、正例と負例の手動設定が必要なため、汎化能力が低いという課題もあります。
  YOLOX は、YOLO シリーズをベースに近年の学術的な物体検出の最新成果を取り入れつつ、YOLO シリーズのデプロイのしやすさを継承しています。さらに、YOLOX では Decoupled Head、Data Aug、Anchor Free、SimOTA の各コンポーネントを設計しました。現在、コードは MegEngine、TensorRT、ONNX、OpenVino、ncnn を含むさまざまなプラットフォームでのデプロイに対応しています。
  EasyCV では、COCO2017 データセット上でさまざまな仕様の事前学習済みモデル (Link) のファインチューニングを提供しており、ダウンストリームタスクに利用できます。
  本記事では、PAI-DSW で EasyCV を使って YOLOX による画像検出モデルのトレーニングと推論を迅速に行う方法を紹介します。
動作環境の要件
PAI-Pytorch 1.8 イメージ、GPU モデル P100 または V100、メモリ 32 GB
依存関係のインストール
1. torch と CUDA のバージョンを取得し、バージョン番号に応じて mmcv のインストールコマンドを変更して、対応バージョンの mmcv と nvidia-dali をインストールします
import torch
import os
os.environ['CUDA']='cu' + torch.version.cuda.replace('.', '')
os.environ['Torch']='torch'+torch.version.__version__.replace('+PAI', '')
。echo $CUDA
。echo $Torch
# install some python deps
。pip install --upgrade tqdm
。pip install mmcv-full==1.4.4 -f https://download.openmmlab.com/mmcv/dist/$CUDA/$Torch/index.html
。pip install http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/third_party/nvidia_dali_cuda100-0.25.0-1535750-py3-none-manylinux2014_x86_64.whl
2. EasyCV アルゴリズムパッケージのインストール 補足：PAI-DSW docker には pai-easycv ライブラリがプリインストール済みなので、この手順はスキップできます
。pip install pai-easycv
# 。pip install http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/pkgs/whl/2022_6/pai_easycv-0.3.0-py3-none-any.whl
3. 簡単な検証
from easycv.apis import *
画像検出モデルのトレーニングと予測
次の例では、cifar10 データを使用し、ResNet50 モデルを使って画像分類モデルのトレーニング評価とモデル予測のプロセスを迅速に実行する方法を紹介します (要修正)
データ準備
COCO2017 データをダウンロードするか、提供されているサンプル COCO データを使用できます
。 wget http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/data/small_coco_demo/small_coco_demo.tar.gz && tar -zxf small_coco_demo.tar.gz
Rename the data file to make it exactly the same as the COCO data format
データファイルをリネームして、COCO データ形式と完全に同じにします
data/coco の形式は以下のとおりです
data/coco/
├── annotations
│ ├── instances_train2017.json
│ └── instances_val2017.json
├── train2017
│ ├── 000000005802.jpg
│ ├── 000000060623.jpg
│ ├── 000000086408.jpg
│ ├── 000000118113.jpg
│ ├── 000000184613.jpg
│ ├── 000000193271.jpg
│ ├── 000000222564.jpg
│...
│ └── 000000574769.jpg
└── val2017
├── 000000006818.jpg
├── 000000017627.jpg
├── 000000037777.jpg
├── 000000087038.jpg
├── 000000174482.jpg
├── 000000181666.jpg
├── 000000184791.jpg
├── 000000252219.jpg
...
モデルのトレーニング
YOLOX-S モデルトレーニング用のサンプル設定ファイルをダウンロードします
Download the sample configuration file for YOLOX-S model training
。rm -rf yolox_s_8xb16_300e_coco.py
小規模データに適応させるため、設定ファイル yolox_s_8xb16_300e_coco.py の次のフィールドを変更して、トレーニングエポック数を減らし、ログ出力の頻度を上げます
In order to adapt to small data, we modify the following fields in the configuration file yolox_s_8xb16_300e_coco.py to reduce the number of training epochs and increase the frequency of printing logs
total_epochs = 3
#optimizer.lr -> 0.0002
optimizer = dict(
type='SGD', lr=0.0002, momentum=0.9, weight_decay=5e-4, nesterov=True)
# log_config.interval 1
補足：COCO の完全データでトレーニングする場合、効果を確保するため 8 GPU 搭載の単一マシンでのトレーニングを推奨します。単一 GPU でトレーニングする場合は、学習率 optimizer.lr を下げることを推奨します
モデルの効果を確保するため、事前学習済みモデルをベースにファインチューニングを行い、次のコマンドを実行してトレーニングを開始します
In order to ensure the effect of the model, we finetune on the basis of the pre-trained model and execute the following command to start the training
モデルのエクスポート
YOLOX モデルを予測用にエクスポートします。次のコマンドを実行して、トレーニングで生成されたモデルファイルを確認します
Export the YOLOX model for prediction, execute the following command to view the model file generated by training
# View the pt file generated by training
モデルをエクスポートする前に、設定ファイルで NMS のスコアしきい値を指定する必要があります
Before exporting the model, the configuration file needs to be modified to specify the score threshold of nms
model.test_conf 0.01 -> 0.5
model = dict(
type='YOLOX',
num_classes=80,
model_type='s', # s m l x tiny nano
test_conf=0.5,
次のコマンドを実行してモデルをエクスポートします
Execute the following command to export the model
。cp yolox_s_8xb16_300e_coco.py yolox_s_8xb16_300e_coco_export.py && sed -i 's#test_conf=0.01#test_conf=0.5#g' yolox_s_8xb16_300e_coco_export.py
予測
テスト画像のダウンロード
Download test image
モデルの重みを読み込み、テスト画像の検出結果を予測します
Import model weights and predict detection results for test images
import cv2
from easycv.predictors import TorchYoloXPredictor
output_ckpt = 'work_dir/detection/yolox/yolox_s_8xb16_300e_coco/yolox_export.pth'
detector = TorchYoloXPredictor(output_ckpt)
img = cv2.imread('000000017627.jpg')
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
output = detector. predict([img])
print(output)
# view detection results
%matplotlib inline
from matplotlib import pyplot as plt
image = img. copy()
for box, cls_name in zip(output[0]['detection_boxes'], output[0]['detection_class_names']):
# box is [x1,y1,x2,y2]
box = [int(b) for b in box]
image = cv2.rectangle(image, tuple(box[:2]), tuple(box[2:4]), (0,255,0), 2)
cv2.putText(image, cls_name, (box[0], box[1]-5), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0,0,255), 2)
plt.imshow(image)
plt. show()