OpenCVのdnnモジュールでYOLOXによるオブジェクト検出を試してみる

この記事は[OpenCV Advent Calendar 2023](https://qiita.com/advent-calendar/2023/opencv)の9日目の記事です。  

# 概要

例年以上に今年の年末は忙しくてネタが思いつきませんでした。
そういえばUnity（Barracuda、Sentis、TensorFlow Lite）では散々実装していたけどOpenCVではやってなかったなとふと思ったので、YOLOXによるオブジェクト検出をOpenCVのdnnモジュールで試してみることにします。

# YOLOX

YOLOXは精度と速度のバランスに優れたオブジェクト検出モデルです。
公式実装のリポジトリも学習、推論、エクスポートの使い勝手がよく整備されています。
商用で扱いやすいライセンスなこともあり、現在も広く利用されている印象です。

* [Megvii-BaseDetection/YOLOX](https://github.com/Megvii-BaseDetection/YOLOX)

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/34834/a86d69c8-11b1-08cf-b487-576ca68eeac8.png)


# サンプルプログラム

サンプルプログラムは以下で公開しています。<sup>*1</sup>　

* [Object Detection using YOLOX inferred by OpenCV DNN module | Gist](https://gist.github.com/UnaNancyOwen/802b724b46977115d94332b5971775b4)

<details><summary>yolox.py</summary><div>

```py
import cv2
import numpy as np

# detection model class for yolox
class DetectionModel:
    # constructor
    def __init__(self, weight, input_size = (640, 640)):
        self.__initialize(weight, input_size)

    # initialize
    def __initialize(self, weight, input_size):
        self.net = cv2.dnn.readNet(weight)
        self.input_size = input_size

        self.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
        self.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)

        strides = [8, 16, 32]
        self.grids, self.expanded_strides = self.__create_grids_and_expanded_strides(strides)

    # create grids and expanded strides
    def __create_grids_and_expanded_strides(self, strides):
        grids = []
        expanded_strides = []

        hsizes = [self.input_size[0] // stride for stride in strides]
        wsizes = [self.input_size[1] // stride for stride in strides]

        for hsize, wsize, stride in zip(hsizes, wsizes, strides):
            xv, yv = np.meshgrid(np.arange(hsize), np.arange(wsize))
            grid = np.stack((xv, yv), 2).reshape(1, -1, 2)
            grids.append(grid)
            shape = grid.shape[:2]
            expanded_strides.append(np.full((*shape, 1), stride))
        
        grids = np.concatenate(grids, 1)
        expanded_strides = np.concatenate(expanded_strides, 1)

        return grids, expanded_strides
    
    # set preferable backend
    def setPreferableBackend(self, backend):
        self.net.setPreferableBackend(backend)
    
    # set preferable target
    def setPreferableTarget(self, target):
        self.net.setPreferableTarget(target)

    # detect objects
    def detect(self, image, score_threshold, iou_threshold):
        self.image_shape = image.shape
        input_blob, resize_ratio = self.__preprocess(image)
        output_blob = self.__predict(input_blob)
        boxes, scores, class_ids = self.__postprocess(output_blob, resize_ratio)
        boxes, scores, class_ids = self.__nms(boxes, scores, class_ids, score_threshold, iou_threshold)

        return class_ids, scores, boxes

    # preprocess
    def __preprocess(self, image):
        resize_ratio = min(self.input_size[0] / self.image_shape[0], self.input_size[1] / self.image_shape[1])
        resized_image = cv2.resize(image, dsize=None, fx=resize_ratio, fy=resize_ratio)

        padded_image = np.ones((self.input_size[0], self.input_size[1], 3), dtype=np.uint8) * 114
        padded_image[: resized_image.shape[0], : resized_image.shape[1]] = resized_image

        input_blob = cv2.dnn.blobFromImage(padded_image, 1.0, self.input_size, (0.0, 0.0, 0.0), True, False)

        return input_blob, resize_ratio
    
    # predict
    def __predict(self, input_blob):
        self.net.setInput(input_blob)

        output_layer = self.net.getUnconnectedOutLayersNames()[0] # "output"
        output_blob = self.net.forward(output_layer)

        return output_blob
    
    # postprocess
    def __postprocess(self, output_blob, resize_ratio):
        output_blob[..., :2] = (output_blob[..., :2] + self.grids) * self.expanded_strides
        output_blob[..., 2:4] = np.exp(output_blob[..., 2:4]) * self.expanded_strides

        predictions = output_blob[0]

        boxes = predictions[:, :4]
        boxes_xywh = np.ones_like(boxes)
        boxes_xywh[:, 0] = boxes[:, 0] - boxes[:, 2] * 0.5
        boxes_xywh[:, 1] = boxes[:, 1] - boxes[:, 3] * 0.5
        boxes_xywh[:, 2] = (boxes[:, 0] + boxes[:, 2] * 0.5) - boxes_xywh[:, 0]
        boxes_xywh[:, 3] = (boxes[:, 1] + boxes[:, 3] * 0.5) - boxes_xywh[:, 1]
        boxes_xywh /= resize_ratio

        scores = predictions[:, 4:5] * predictions[:, 5:]
        class_ids = scores.argmax(1)
        scores = scores[np.arange(len(class_ids)), class_ids]

        return boxes_xywh, scores, class_ids
    
    # non maximum suppression
    def __nms(self, boxes, scores, class_ids, score_threshold, iou_threshold):
        indices = cv2.dnn.NMSBoxesBatched(boxes, scores, class_ids, score_threshold, iou_threshold) # OpenCV 4.7.0 or later

        keep_boxes = []
        keep_scores = []
        keep_class_ids = []
        for index in indices:
            keep_boxes.append(boxes[index])
            keep_scores.append(scores[index])
            keep_class_ids.append(class_ids[index])

        if len(keep_boxes) != 0:
            keep_boxes = np.vectorize(int)(keep_boxes)

        return keep_boxes, keep_scores, keep_class_ids
```
</div></details>

# 使い方

OpenCVのdnnモジュールのHight Level APIっぽいインターフェースでクラスにまとめました。
中身が知りたい人は[yolox.py](https://gist.github.com/UnaNancyOwen/802b724b46977115d94332b5971775b4#file-yolox-py)を読んでください。ここでは使い方を掻い摘んでざっくり説明します。

## ・インポート

yolox.pyで実装されているDetectionModelクラスをインポートします。  
このクラスでYOLOXの推論機能を提供します。

```py
from yolox import DetectionModel
```

## ・クラスのインスタンス生成

クラスのインスタンスを生成します。
学習済みモデルのONNXファイルのパスのほかに、入力レイヤーの形状（幅、高さ）を渡します。
OpenCVのdnnモジュールで入力レイヤーの形状を取得できないのでこのような仕様になっています。（誰か教えて

```py
# create detection model class for yolox
weight = "yolox_s.onnx"
input_size = (640, 640)
model = DetectionModel(weight, input_size)
```

## ・オブジェクトを検出

`detect()`メソッドに画像、信頼度の閾値、IoUの閾値を渡してオブジェクトを検出します。
戻り値はクラスID、信頼度、バウンディングボックスのリストです。
バウンディングボックスは(x, y, width, height)の形式です。

```py
# detect objects
score_threshold = 0.6
iou_threshold = 0.4
class_ids, scores, boxes = model.detect(image, score_threshold, iou_threshold)
```

## ・オブジェクトを描画

検出したオブジェクトを描画します。
バウンディングボックスはOpenCVの[Rect](https://docs.opencv.org/4.8.0/d2/d44/classcv_1_1Rect__.html)と同じ形式になっています。
そのため、そのまま[cv.rectangle()](https://docs.opencv.org/4.8.0/d6/d6e/group__imgproc__draw.html#gac865734d137287c0afb7682ff7b3db23)に渡して描画することができます。

```py
for box, score, class_id in zip(boxes, scores, class_ids):
    color = (0, 0, 255)
    thickness = 2
    line_type = cv2.LINE_AA
    cv2.rectangle(image, box, color, thickness, line_type)
```

# 結果

* [Official YOLOX Pre-Trained Model](https://github.com/Megvii-BaseDetection/YOLOX/tree/main/demo/ONNXRuntime#download-onnx-models) <sup>*2</sup>
    YOLOXの公式リポジトリで配布されている学習済みモデルで試してみた結果です。
    ![yolox_s.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/34834/17633fdb-da40-1bc6-1578-02c865c4ceea.png)

* [PINTO_model_zoo YOLOX Pre-Trained Model for Body/Head/Hand](https://github.com/PINTO0309/PINTO_model_zoo/tree/main/426_YOLOX-Body-Head-Hand) <sup>*3 *4</sup>
    [PINTO](https://qiita.com/PINTO)さんが作成したYOLOXで身体、頭、手を検出する学習済みモデルで試してみた結果です。
    ![image6.jpg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/34834/940d2085-fe9c-caec-678d-76bf7cdfea11.jpeg)

<!--![yolox_s_body_head_hand.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/34834/f401b660-0639-2a95-7f32-567199df1a73.png)-->

# まとめ

この記事ではOpenCVのdnnモジュールでYOLOXによるオブジェクト検出を試してみました。
公式のONNX Runtimeでの実装とあまり変わり映えしないけど、OpenCV 4.7.0あたりで追加されたマルチクラス対応のNMSのAPI（`cv.dnn.NMSBoxesBatched()`）が用意されているのは少し楽かもしれない。

---

<sup>*1 ちゃちゃっと実装したのでモデルの入力形状のアスペクト比が1:1でないと上手く動かないかも</sup>
<sup>*2 画像引用：https://github.com/Megvii-BaseDetection/YOLOX/blob/main/assets/dog.jpg</sup>
<sup>*3 画像引用：https://www.flickr.com/photos/83863691@N00/3704144136</sup>
<sup>*4 ポストプロセスが統合されていないONNXを使ってください</sup>