CVATにカスタムYOLOモデルをデプロイする

今回はCVATでYOLOのカスタムモデルをデプロイする方法を書いていきます。CVATは無料で使えるアノテーションツールですが、AIモデルを載せることで画像の分析基盤としても活用することができます。

今回のこの記事の内容は以下の内容を参考にしています。

https://medium.com/@alhasan.saad/automating-data-annotation-how-to-set-up-cvat-and-nuclio-for-custom-and-built-in-models-ac49e3d4d931

https://medium.com/@eng.fadishaar/automating-object-annotation-in-cvat-using-a-custom-yolov5-model-cfd36fb40a97

## 前提

この記事はCVATの設定が終わっていて、CVATが用意しているAIモデルでの自動アノテーションが使える状態になった後を想定しています。ここまでの準備については以下の記事を参照してください。

https://qiita.com/hbk24/items/7eff8183b98574d4ad56

まずは稼働中のCVATのDockerコンテナを停止させます。

```shell
cd ~/cvat
docker compose -f docker-compose.yml -f components/serverless/docker-compose.serverless.yml down
```

## サーバーレス関数の実装

CVATでは`Nuclio`を使ってAIのサーバーレス関数を構築し、CVATからこれを呼び出すことでAIによる処理を行っています。この項では、`Nuclio`へのサーバーレス関数の実装を解説します。

https://nuclio.io/

また、今回は`Ultralytics`のYOLOv9をデプロイしていきます。

https://docs.ultralytics.com/models/yolov9/

`Ultralytics`を使えば重みファイルをわざわざダウンロードしなくても大丈夫ですが、カスタムしたモデルをデプロイする場合は重みファイルを指定しなくてはなりません。そこで、デフォルトモデルの重みファイルをダウンロードして使用し、モデルのカスタムが終わったら重みファイルを差し替える想定で進めます。

まず、デフォルトの重みファイル（yolov9c.py）をダウンロードします。
https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov9c.pt

次にCVATフォルダ内にあるserverlessフォルダ内へ今回デプロイする用のフォルダを作ります。

``` shell
mkdir -p serverless/ultralytics/yolov9/nuclio
```

ちなみに、フォルダの階層やフォルダ名は自由に変更してもらっても問題ありません。単純に、フレームワークを区別するために、このような階層構造にしています。

次にこのフォルダ内に以下のファイルを置きます。

- yolov9c.pt　(カスタムした場合はその重みファイルを置く)
- main.py
- function.yaml (AIをCPUで処理する場合)
- function-gpu.yaml (AIをCPUで処理する場合)

各フォルダには以下のコードを書いていきます。

### YAML

YAMLファイルでは`Nuclio`のサーバーレス関数を実行するための設定を記述しています。`DockerFile`のようなイメージです。コードの詳細は以下のコードに付けているコメントを参照してください。

ちなみに、下で解説している通り、`Nuclio`はGUIのダッシュボードが用意されいて、そこで設定をすることも可能です。ダッシュボードを覗いてみるとコードの中身が理解しやすくなると思います。

```yaml:function-gpu.yaml
metadata:
  name: gpu-ultralytics-yolov9 # Nuclio内での表示名
  namespace: cvat # Nuclioのプロジェクト名
  annotations:
    name: GPU YOLOv9 by Ultralytics # CVATで表示されるモデル名
    type: detector
    framework: pytorch
    # モデルが識別するラベル名のリスト（カスタムモデルの場合、それに合わせる）
    spec: |
      [
        { "id": 0, "name": "person", "type": "rectangle" },
        { "id": 1, "name": "bicycle", "type": "rectangle" },
        { "id": 2, "name": "car", "type": "rectangle" },
        { "id": 3, "name": "motorbike", "type": "rectangle" },
        { "id": 4, "name": "aeroplane", "type": "rectangle" },
        { "id": 5, "name": "bus", "type": "rectangle" },
        { "id": 6, "name": "train", "type": "rectangle" },
        { "id": 7, "name": "truck", "type": "rectangle" },
        { "id": 8, "name": "boat", "type": "rectangle" },
        { "id": 9, "name": "traffic light", "type": "rectangle" },
        { "id": 10, "name": "fire hydrant", "type": "rectangle" },
        { "id": 11, "name": "stop sign", "type": "rectangle" },
        { "id": 12, "name": "parking meter", "type": "rectangle" },
        { "id": 13, "name": "bench", "type": "rectangle" },
        { "id": 14, "name": "bird", "type": "rectangle" },
        { "id": 15, "name": "cat", "type": "rectangle" },
        { "id": 16, "name": "dog", "type": "rectangle" },
        { "id": 17, "name": "horse", "type": "rectangle" },
        { "id": 18, "name": "sheep", "type": "rectangle" },
        { "id": 19, "name": "cow", "type": "rectangle" },
        { "id": 20, "name": "elephant", "type": "rectangle" },
        { "id": 21, "name": "bear", "type": "rectangle" },
        { "id": 22, "name": "zebra", "type": "rectangle" },
        { "id": 23, "name": "giraffe", "type": "rectangle" },
        { "id": 24, "name": "backpack", "type": "rectangle" },
        { "id": 25, "name": "umbrella", "type": "rectangle" },
        { "id": 26, "name": "handbag", "type": "rectangle" },
        { "id": 27, "name": "tie", "type": "rectangle" },
        { "id": 28, "name": "suitcase", "type": "rectangle" },
        { "id": 29, "name": "frisbee", "type": "rectangle" },
        { "id": 30, "name": "skis", "type": "rectangle" },
        { "id": 31, "name": "snowboard", "type": "rectangle" },
        { "id": 32, "name": "sports ball", "type": "rectangle" },
        { "id": 33, "name": "kite", "type": "rectangle" },
        { "id": 34, "name": "baseball bat", "type": "rectangle" },
        { "id": 35, "name": "baseball glove", "type": "rectangle" },
        { "id": 36, "name": "skateboard", "type": "rectangle" },
        { "id": 37, "name": "surfboard", "type": "rectangle" },
        { "id": 38, "name": "tennis racket", "type": "rectangle" },
        { "id": 39, "name": "bottle", "type": "rectangle" },
        { "id": 40, "name": "wine glass", "type": "rectangle" },
        { "id": 41, "name": "cup", "type": "rectangle" },
        { "id": 42, "name": "fork", "type": "rectangle" },
        { "id": 43, "name": "knife", "type": "rectangle" },
        { "id": 44, "name": "spoon", "type": "rectangle" },
        { "id": 45, "name": "bowl", "type": "rectangle" },
        { "id": 46, "name": "banana", "type": "rectangle" },
        { "id": 47, "name": "apple", "type": "rectangle" },
        { "id": 48, "name": "sandwich", "type": "rectangle" },
        { "id": 49, "name": "orange", "type": "rectangle" },
        { "id": 50, "name": "broccoli", "type": "rectangle" },
        { "id": 51, "name": "carrot", "type": "rectangle" },
        { "id": 52, "name": "hot dog", "type": "rectangle" },
        { "id": 53, "name": "pizza", "type": "rectangle" },
        { "id": 54, "name": "donut", "type": "rectangle" },
        { "id": 55, "name": "cake", "type": "rectangle" },
        { "id": 56, "name": "chair", "type": "rectangle" },
        { "id": 57, "name": "sofa", "type": "rectangle" },
        { "id": 58, "name": "pottedplant", "type": "rectangle" },
        { "id": 59, "name": "bed", "type": "rectangle" },
        { "id": 60, "name": "diningtable", "type": "rectangle" },
        { "id": 61, "name": "toilet", "type": "rectangle" },
        { "id": 62, "name": "tvmonitor", "type": "rectangle" },
        { "id": 63, "name": "laptop", "type": "rectangle" },
        { "id": 64, "name": "mouse", "type": "rectangle" },
        { "id": 65, "name": "remote", "type": "rectangle" },
        { "id": 66, "name": "keyboard", "type": "rectangle" },
        { "id": 67, "name": "cell phone", "type": "rectangle" },
        { "id": 68, "name": "microwave", "type": "rectangle" },
        { "id": 69, "name": "oven", "type": "rectangle" },
        { "id": 70, "name": "toaster", "type": "rectangle" },
        { "id": 71, "name": "sink", "type": "rectangle" },
        { "id": 72, "name": "refrigerator", "type": "rectangle" },
        { "id": 73, "name": "book", "type": "rectangle" },
        { "id": 74, "name": "clock", "type": "rectangle" },
        { "id": 75, "name": "vase", "type": "rectangle" },
        { "id": 76, "name": "scissors", "type": "rectangle" },
        { "id": 77, "name": "teddy bear", "type": "rectangle" },
        { "id": 78, "name": "hair drier", "type": "rectangle" },
        { "id": 79, "name": "toothbrush", "type": "rectangle" }
      ]

spec:
  description: GPU YOLOv9 by Ultralytics
  runtime: 'python:3.9' # サーバーレスファンクションを実行するランタイムを指定（Nuclio公式は3.9以上を推奨）
  handler: main:handler # トリガーを引かれた時に実行するプログラムファイル名（ここでは"main.py"を想定）
  eventTimeout: 30s

  build:
    image: cvat.yolov9.gpu # 構成するDockerイメージ名
    baseImage: ultralytics/ultralytics # Dockerのベース

    # 以下DockerFileのような部分
    directives:
      preCopy:
        - kind: USER
          value: root
        # NVIDIAコンテナのセッティング
        - kind: ENV
          value: NVIDIA_VISIBLE_DEVICES=all
        - kind: ENV
          value: NVIDIA_DRIVER_CAPABILITIES=compute,utility
        - kind: ENV
          value: RUN_ON_GPU="true" #cpuの場合”false”
        # Pythonの設定
        - kind: RUN
          value: export DEBIAN_FRONTEND=noninteractive && apt-get update && apt-get install -y python-is-python3
        - kind: RUN
          value: pip install --no-cache-dir opencv-python-headless pillow pyyaml
        - kind: RUN
          value: pip uninstall -y torch torchvision torchaudio
        - kind: RUN
          value: pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
        - kind: WORKDIR
          value: /opt/nuclio

  # トリガーの設定
  triggers:
    myHttpTrigger:
      maxWorkers: 1
      kind: 'http'
      workerAvailabilityTimeoutMilliseconds: 10000
      attributes:
        maxRequestBodySize: 33554432 # 32MB

  # GPUの設定（最大何個使用するか）CPUの場合はここを消す
  resources:
    limits:
      nvidia.com/gpu: 1

  # サーバーレスファンクションの追加設定
  platform:
    attributes:
      restartPolicy:
        name: always
        maximumRetryCount: 3
      mountMode: volume
```

```yaml:function.yaml
metadata:
  name: cpu-ultralytics-yolov9 # Nuclio内での表示名
  namespace: cvat # Nuclioのプロジェクト名
  annotations:
    name: CPU YOLOv9 by Ultralytics # CVATで表示されるモデル名
    type: detector
    framework: pytorch
    # モデルが識別するラベル名のリスト（カスタムモデルの場合、それに合わせる）
    spec: |
      [
        { "id": 0, "name": "person", "type": "rectangle" },
        { "id": 1, "name": "bicycle", "type": "rectangle" },
        { "id": 2, "name": "car", "type": "rectangle" },
        { "id": 3, "name": "motorbike", "type": "rectangle" },
        { "id": 4, "name": "aeroplane", "type": "rectangle" },
        { "id": 5, "name": "bus", "type": "rectangle" },
        { "id": 6, "name": "train", "type": "rectangle" },
        { "id": 7, "name": "truck", "type": "rectangle" },
        { "id": 8, "name": "boat", "type": "rectangle" },
        { "id": 9, "name": "traffic light", "type": "rectangle" },
        { "id": 10, "name": "fire hydrant", "type": "rectangle" },
        { "id": 11, "name": "stop sign", "type": "rectangle" },
        { "id": 12, "name": "parking meter", "type": "rectangle" },
        { "id": 13, "name": "bench", "type": "rectangle" },
        { "id": 14, "name": "bird", "type": "rectangle" },
        { "id": 15, "name": "cat", "type": "rectangle" },
        { "id": 16, "name": "dog", "type": "rectangle" },
        { "id": 17, "name": "horse", "type": "rectangle" },
        { "id": 18, "name": "sheep", "type": "rectangle" },
        { "id": 19, "name": "cow", "type": "rectangle" },
        { "id": 20, "name": "elephant", "type": "rectangle" },
        { "id": 21, "name": "bear", "type": "rectangle" },
        { "id": 22, "name": "zebra", "type": "rectangle" },
        { "id": 23, "name": "giraffe", "type": "rectangle" },
        { "id": 24, "name": "backpack", "type": "rectangle" },
        { "id": 25, "name": "umbrella", "type": "rectangle" },
        { "id": 26, "name": "handbag", "type": "rectangle" },
        { "id": 27, "name": "tie", "type": "rectangle" },
        { "id": 28, "name": "suitcase", "type": "rectangle" },
        { "id": 29, "name": "frisbee", "type": "rectangle" },
        { "id": 30, "name": "skis", "type": "rectangle" },
        { "id": 31, "name": "snowboard", "type": "rectangle" },
        { "id": 32, "name": "sports ball", "type": "rectangle" },
        { "id": 33, "name": "kite", "type": "rectangle" },
        { "id": 34, "name": "baseball bat", "type": "rectangle" },
        { "id": 35, "name": "baseball glove", "type": "rectangle" },
        { "id": 36, "name": "skateboard", "type": "rectangle" },
        { "id": 37, "name": "surfboard", "type": "rectangle" },
        { "id": 38, "name": "tennis racket", "type": "rectangle" },
        { "id": 39, "name": "bottle", "type": "rectangle" },
        { "id": 40, "name": "wine glass", "type": "rectangle" },
        { "id": 41, "name": "cup", "type": "rectangle" },
        { "id": 42, "name": "fork", "type": "rectangle" },
        { "id": 43, "name": "knife", "type": "rectangle" },
        { "id": 44, "name": "spoon", "type": "rectangle" },
        { "id": 45, "name": "bowl", "type": "rectangle" },
        { "id": 46, "name": "banana", "type": "rectangle" },
        { "id": 47, "name": "apple", "type": "rectangle" },
        { "id": 48, "name": "sandwich", "type": "rectangle" },
        { "id": 49, "name": "orange", "type": "rectangle" },
        { "id": 50, "name": "broccoli", "type": "rectangle" },
        { "id": 51, "name": "carrot", "type": "rectangle" },
        { "id": 52, "name": "hot dog", "type": "rectangle" },
        { "id": 53, "name": "pizza", "type": "rectangle" },
        { "id": 54, "name": "donut", "type": "rectangle" },
        { "id": 55, "name": "cake", "type": "rectangle" },
        { "id": 56, "name": "chair", "type": "rectangle" },
        { "id": 57, "name": "sofa", "type": "rectangle" },
        { "id": 58, "name": "pottedplant", "type": "rectangle" },
        { "id": 59, "name": "bed", "type": "rectangle" },
        { "id": 60, "name": "diningtable", "type": "rectangle" },
        { "id": 61, "name": "toilet", "type": "rectangle" },
        { "id": 62, "name": "tvmonitor", "type": "rectangle" },
        { "id": 63, "name": "laptop", "type": "rectangle" },
        { "id": 64, "name": "mouse", "type": "rectangle" },
        { "id": 65, "name": "remote", "type": "rectangle" },
        { "id": 66, "name": "keyboard", "type": "rectangle" },
        { "id": 67, "name": "cell phone", "type": "rectangle" },
        { "id": 68, "name": "microwave", "type": "rectangle" },
        { "id": 69, "name": "oven", "type": "rectangle" },
        { "id": 70, "name": "toaster", "type": "rectangle" },
        { "id": 71, "name": "sink", "type": "rectangle" },
        { "id": 72, "name": "refrigerator", "type": "rectangle" },
        { "id": 73, "name": "book", "type": "rectangle" },
        { "id": 74, "name": "clock", "type": "rectangle" },
        { "id": 75, "name": "vase", "type": "rectangle" },
        { "id": 76, "name": "scissors", "type": "rectangle" },
        { "id": 77, "name": "teddy bear", "type": "rectangle" },
        { "id": 78, "name": "hair drier", "type": "rectangle" },
        { "id": 79, "name": "toothbrush", "type": "rectangle" }
      ]

spec:
  description: CPU YOLOv9 by Ultralytics
  runtime: 'python:3.9' # サーバーレスファンクションを実行するランタイムを指定（Nuclio公式は3.9以上を推奨）
  handler: main:handler # トリガーを引かれた時に実行するプログラムファイル名（ここでは"main.py"を想定）
  eventTimeout: 30s

  build:
    image: cvat.yolov9.cpu # 構成するDockerイメージ名
    baseImage: ultralytics/ultralytics # Dockerのベース

    # 以下DockerFileのような部分
    directives:
      preCopy:
        - kind: USER
          value: root
        # NVIDIAコンテナのセッティング
        - kind: ENV
          value: NVIDIA_VISIBLE_DEVICES=all
        - kind: ENV
          value: NVIDIA_DRIVER_CAPABILITIES=compute,utility
        - kind: ENV
          value: RUN_ON_GPU="false"
        # Pythonの設定
        - kind: RUN
          value: export DEBIAN_FRONTEND=noninteractive && apt-get update && apt-get install -y python-is-python3
        - kind: RUN
          value: pip install --no-cache-dir opencv-python-headless pillow pyyaml
        - kind: RUN
          value: pip uninstall -y torch torchvision torchaudio
        - kind: RUN
          value: pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
        - kind: WORKDIR
          value: /opt/nuclio

  # トリガーの設定
  triggers:
    myHttpTrigger:
      maxWorkers: 1
      kind: 'http'
      workerAvailabilityTimeoutMilliseconds: 10000
      attributes:
        maxRequestBodySize: 33554432 # 32MB

  # サーバーレスファンクションの追加設定
  platform:
    attributes:
      restartPolicy:
        name: always
        maximumRetryCount: 3
      mountMode: volume
```

### main.py

`init_context`: サーバーレス関数のコンテキストを初期化を行っています。YOLOモデルを読み込みや、GPUの使用が有効か確認を行っています。

`handler()`: 処理のメイン部分です・
- 入力処理：画像データを含むHTTPリクエストによってトリガーされます。画像データを抽出してデコードし、オブジェクト検出の信頼度閾値を設定し、Pillowを使用して画像を読み込んでいます。
- オブジェクト検出：読み込まれたモデルは、設定された閾値に基づいて画像からオブジェクトを検出します。context.user_dataに保存されているモデル（）が検出を実行し、結果を返します。
- 結果のフォーマット：検出されたオブジェクトは、CVATが理解できるJSON構造にフォーマットされます。スクリプトは、検出された各オブジェクトについて、信頼度スコア、クラスラベル、バウンディングボックスの座標を抽出しています。
- HTTPレスポンス：関数は、JSON形式の検出結果を含むHTTPレスポンスを構築します。レスポンスにはヘッダーとコンテンツタイプの仕様が含まれ、ステータスコード200は処理が成功したことを示しています。

```python:main.py
from ultralytics import YOLO
import json, base64, io, os
from PIL import Image

def init_context(context):
    context.logger.info("Init context...  0%")

    # main.pyと同じフォルダにあるモデルの重みファイル（’yolov9c.pt’）を読み込み
    model = YOLO('/opt/nuclio/yolov9c.pt')

    use_gpu = os.getenv("RUN_ON_GPU", 'False').lower() in ('true', '1') # Import the GPU env variable and covert to a boolean value
    print(f"CUDA-STATUS: {use_gpu}")
    if use_gpu:
        model.to('cuda')

    context.user_data.model = model

    context.logger.info("Init context...100%")

def handler(context, event):
    context.logger.info("Run yolo-v9 model")
    data = event.body
    buf = io.BytesIO(base64.b64decode(data["image"]))
    threshold = float(data.get("threshold", 0.6))
    context.user_data.model.conf = threshold
    image = Image.open(buf)
    yolo_results = context.user_data.model(image, conf=threshold)
    results = yolo_results[0]

    encoded_results = []
    for idx, class_idx in enumerate(results.boxes.cls):
        confidence = results.boxes.conf[idx].item()
        label = results.names[int(class_idx.item())]
        points = results.boxes.xyxy[idx].tolist()
        encoded_results.append({
            'confidence': confidence,
            'label': label,
            'points': points,
            'type': 'rectangle'
        })

    return context.Response(body=json.dumps(encoded_results), headers={},
        content_type='application/json', status_code=200)
```

### デプロイ

GPUを使用する場合は

```shell
./serverless/deploy_gpu.sh ./serverless/ultralytics/yolov9/nuclio/
```

CPUの場合は

```shell
./serverless/deploy_gpu.sh ./serverless/ultralytics/yolov9/nuclio/
```

を実行し、しばらくすると


```shell
 NAMESPACE | NAME                           | PROJECT | STATE | REPLICAS | NODE PORT   
 nuclio    | gpu-ultralytics-yolov9         | cvat    | ready | 1/1      | 33311   
```

とターミナルで返されるのでこれでデプロイ完了です。

## CVATでカスタムYOLOモデルを使用

CVATをサーバーPCに入れている場合は

```shell
export CVAT_HOST=<IPアドレス>
```

を実行してから

```shell
docker compose -f docker-compose.yml -f components/serverless/docker-compose.serverless.yml up -d
```

でCVATのコンテナを起動させます。

アノテーション画面で「AI Tools」を開くとデプロイしたYOLOのモデルが使用できるようになっています。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/773995/72779823-779e-b22d-a116-bbbce41fdcc8.png)
(画像はCOCOのデータセットに含まれる画像です)

## サーバーレス関数の削除

実装したサーバーレス関数を削除したい場合はCVATのコンテナが走っている状態で

[http://IPアドレス/:8070]()

にアクセスすると`Nuclio`のダッシュボードにアクセスできます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/773995/90f81f3b-4855-7c2e-fa9f-4bc7f22af09a.png)

ここから、実装したサーバーレス関数を削除することができます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/773995/77c322ad-b9b4-2cfd-d442-bb7ac0920654.png)

CVATのコンテナが走っていない状態では、以下のコマンドで`Nuclio`のダッシュボードのコンテナを立ち上げるとダッシュボードにアクセスできます。

```shell
docker run --detach --publish 8070:8070 --volume /var/run/docker.sock:/var/run/docker.sock --name nuclio-dashboard quay.io/nuclio/dashboard:stable-amd64
```
https://docs.nuclio.io/en/stable/tasks/quick-start.html