RT-DETRを自作データで学習するフローをまとめました

# はじめに

皆さんはRT-DETRというモデルを知っていますか？ YOLOv8よりも高速で高精度なTransformerベースのObject Detectionモデルです。なんかすごそう(小並感)

詳しくは以下を参照してください。

https://arxiv.org/abs/2304.08069

![238573637-3184a08e-aa4d-49cf-9079-f3695c4cc1c3.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/899689/2168561f-9ee7-c047-54c6-a6fc6f4af008.png)

# 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環境でも実行可能です。その方は以下のサイトを参考に環境構築してください。

https://docs.ultralytics.com/quickstart/#use-ultralytics-with-python


# 1. データセットの用意
お手持ちのスマホやカメラで撮影した画像データをたくさん用意してください。自分はテスト用にグー、チョキ、パーの手の画像を用意しました。ただしpngやjpgの形式にするのを忘れないでください。

そして今回は私が作った†楽々最強AIツール†を使っていきます。

まずは`git clone`

```bash
git clone https://github.com/TakanoTaiga/ml.git
```

クローンするとこんなファイルやフォルダが入ったリポジトリがクローンされます。

![Screenshot 2024-04-02 at 21.39.58.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/899689/99534135-00fd-6272-b6b5-f1caa3e10427.png)

そしたら先ほど用意した画像たちを`input_image`フォルダにすべて入れてください。(この時何か事前の加工等は必要ありません)

次にターミナルを開いて画像の大きさを整えすべてjpgにしてファイル名を整えるスクリプトを実行します。

```bash
cd ./ml
python3 set_format.py
```

そうすると`out_image`フォルダが生成されていることがわかると思います。中を確認して画像がたくさんあればOKです。

![Screenshot 2024-04-02 at 21.44.31.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/899689/736dd822-4c0f-75b4-c0e9-bc16b61921ee.png)


# アノテーション
アノテーションツールはcoco-annotatorを使います。

https://github.com/jsbroks/coco-annotator

使い方は簡単`docekr-compose`でポン。

```bash
git clone https://github.com/jsbroks/coco-annotator.git
cd coco-annotator
docekr-compose up
```
これで [localhost:5000](http://localhost:5000/)にアクセスすればOKです。
ctrl-cで終了できますがその後は下記のコマンドで起動と停止を行なってください。

```bash:start
cd coco-annotator
docker-compose start
```

```bash:stop
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`ファイルが生成されて入れが完了です。

![Screenshot 2024-04-02 at 22.01.56.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/899689/83bc0ca7-1d16-20b9-c74b-c140485be942.png)


いよいよ学習を回していきます。

**jetson環境**
コンテナを起動し先ほど生成されたpythonファイルを実行します。
```bash
./start.sh
```
起動後(xxxxxの部分は適宜変えてください.)
```bash
python3 train_rtdetr_xxxxx.py
```

**x86+NVIDIA GPU環境**
コンテナを起動することなく先ほど生成されたpythonファイルを実行します。
```bash
python3 train_rtdetr_xxxxx.py
```

学習が完了すると`ml/run/trainNN`の中に重みファイルやログデータが生成されます。これで学習完了です。


# 推論

`train/weight`フォルダの中に`.pt`ファイルがあると思うのでbestかlastお好きな方を適当なフォルダに置きます。そうしたら下記のコードを適当なpythonファイルに保存して実行します。カメラからの映像を自動的に読み取って推論が始まります。

`hogehoge.pt`の部分は置き換えてください。

```python
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()
```