はじめに
皆さんはRT-DETRというモデルを知っていますか？ YOLOv8よりも高速で高精度なTransformerベースのObject Detectionモデルです。なんかすごそう(小並感)
詳しくは以下を参照してください。
0. 環境構築(ハードウェア編)
RT-DETRの学習にはlargeモデルで13GBくらいのVRAMを使用しますので16GBクラスのVRAMが搭載されたGPUが必要になります。ここでApple Silicon Macは？という声があるかと思いますが、M1 MAXがRTX3050相当のGPUの3倍くらいの時間をかけて学習していたので使うのは厳しいかと思います。
今回おすすめなのがJetson AGX Orinです。CUDAやcuDNNが最初からセットアップされているので非常におすすめです。
x86+NVIDIA GPU環境でも実行可能です。その方は以下のサイトを参考に環境構築してください。
1. データセットの用意
お手持ちのスマホやカメラで撮影した画像データをたくさん用意してください。自分はテスト用にグー、チョキ、パーの手の画像を用意しました。ただしpngやjpgの形式にするのを忘れないでください。
そして今回は私が作った†楽々最強AIツール†を使っていきます。
まずはgit clone
git clone https://github.com/TakanoTaiga/ml.git
クローンするとこんなファイルやフォルダが入ったリポジトリがクローンされます。
そしたら先ほど用意した画像たちをinput_imageフォルダにすべて入れてください。(この時何か事前の加工等は必要ありません)
次にターミナルを開いて画像の大きさを整えすべてjpgにしてファイル名を整えるスクリプトを実行します。
cd ./ml
python3 set_format.py
そうするとout_imageフォルダが生成されていることがわかると思います。中を確認して画像がたくさんあればOKです。
アノテーション
アノテーションツールはcoco-annotatorを使います。
使い方は簡単docekr-composeでポン。
git clone https://github.com/jsbroks/coco-annotator.git
cd coco-annotator
docekr-compose up
これで localhost:5000にアクセスすればOKです。
ctrl-cで終了できますがその後は下記のコマンドで起動と停止を行なってください。
cd coco-annotator
docker-compose start
cd coco-annotator
docker-compose stop
起動できたらcoco-annotatorでdatasetを作り先ほどのout_imageフォルダの中身をコピーします。(この時out_image内の画像を削除したり切り取りコピーしないでください。)
アノテーションが完了したアノテーション結果のjsonファイルをダウンロードしてください。そのjsonファイルはml/input_labelフォルダにコピーします。
学習
データが完成したら次は学習です。学習するためには現在のcocoフォーマットからyolo形式に変換する必要があります。そのためのスクリプトを用意したので実行します。
cd ml
python3 coco2yolo.py
実行するとたくさんのフォルダが生成されます。train_rtdetr_xxxxx.pyファイルが生成されて入れが完了です。
いよいよ学習を回していきます。
jetson環境
コンテナを起動し先ほど生成されたpythonファイルを実行します。
./start.sh
起動後(xxxxxの部分は適宜変えてください.)
python3 train_rtdetr_xxxxx.py
x86+NVIDIA GPU環境
コンテナを起動することなく先ほど生成されたpythonファイルを実行します。
python3 train_rtdetr_xxxxx.py
学習が完了するとml/run/trainNNの中に重みファイルやログデータが生成されます。これで学習完了です。
推論
train/weightフォルダの中に.ptファイルがあると思うのでbestかlastお好きな方を適当なフォルダに置きます。そうしたら下記のコードを適当なpythonファイルに保存して実行します。カメラからの映像を自動的に読み取って推論が始まります。
hogehoge.ptの部分は置き換えてください。
import cv2
from ultralytics import RTDETR
model = RTDETR('hogehoge.pt')
# Open the web camera stream
cap = cv2.VideoCapture(0)
while cap.isOpened():
success, frame = cap.read()
k = cv2.waitKey(1)
if k != -1:
break
if success:
results = model.predict(frame,conf=0.7,half=True)
annotated_frame = results[0].plot()
cv2.imshow("RT-DETR Inference", annotated_frame)
cap.release()
cv2.destroyAllWindows()