OpenCV의 DNN 모듈

OpenCV 4.1.2에서는 DNN 모듈에 CUDA 옵션이 추가되었습니다. 이 DNN 모듈은 다양한 프레임워크에서 생성된 학습된 모델을 읽고 실행할 수 있습니다.
일반 물체 인식의 고속 모델로서 YOLOv3가 있습니다만, 이것도 readNetFromDarknet 함수로 읽어들여 추론을 실시할 수 있습니다.

main.cpp

#include <opencv2/core.hpp>
#include <opencv2/dnn.hpp>
#include <opencv2/highgui.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/videoio.hpp>

using namespace cv;
using namespace dnn;

/// Darknetのモデル
static const char* cfg = "yolov3.cfg";
static const char* weights = "yolov3.weights";

/// 推論結果のヒートマップを解析して表示
void postprocess(Mat& frame, const std::vector<Mat>& outs, Net& net);

static float confThreshold = 0.5f;
static float nmsThreshold = 0.4f;
std::vector<std::string> classes;

int main(int argc, char *argv[])
{
    VideoCapture camera;
    if (camera.open(0))
    {
        /// YOLOv3のモデルを読み込む
        Net net = readNetFromDarknet(cfg, weights);
        // OpenCV 4.1.2をソースからCUDAを有効にしてビルドした場合、このようにCUDAを指定できます
        net.setPreferableBackend(DNN_BACKEND_CUDA);
        net.setPreferableTarget(DNN_TARGET_CUDA);

        Mat image, blob;
        std::vector<Mat> outs;
        std::vector<String> outNames = net.getUnconnectedOutLayersNames();
        for (bool loop = true; loop;)
        {
            camera >> image;

            /// 画像をBLOBに変換して推論
            blob = blobFromImage(image, 1.0f / 255);
            net.setInput(blob);
            net.forward(outs, outNames);

            /// 推論結果をimageに描画
            postprocess(image, outs, net);

            imshow("Image", image);
            switch (waitKey(10))
            {
            case 'q':
            case 'Q':
                loop = false;
            }
        }
        camera.release();
    }
}

추론 결과는 히트맵이므로 여기에서 필요한 정보를 추출합니다. 이 부분은 모델에 따라 다르지만 YOLOv3에서는 영역 데이터 뒤에 범주를 나타내는 히트 맵이 뒤따르고 있으며 범주 신뢰도를 postprocess 함수로 추출하고 그립니다.

void postprocess(Mat& frame, const std::vector<Mat>& outs, Net& net)
{
    static std::vector<int> outLayers = net.getUnconnectedOutLayers();
    static std::string outLayerType = net.getLayer(outLayers[0])->type;

    if (outLayerType == "Region")
    {
        for (Mat out : outs)
        {
            float* data = (float*)out.data;

            // 検出したオブジェクトごとに
            for (int i = 0; i < out.rows; i++, data += out.cols)
            {
                // 領域情報
                int centerX = (int)(data[0] * frame.cols);
                int centerY = (int)(data[1] * frame.rows);
                int width = (int)(data[2] * frame.cols);
                int height = (int)(data[3] * frame.rows);

                // そのあとに一次元のヒートマップが続く
                Mat scores = out.row(i).colRange(5, out.cols);
                Point classIdPoint;
                double confidence;

                // 信頼度とクラスを抽出
                minMaxLoc(scores, 0, &confidence, 0, &classIdPoint);

                // 信頼度が閾値超えたオブジェクトを描画する
                if (confThreshold < confidence)
                {
                    int left = centerX - width / 2;
                    int top = centerY - height / 2;
                    // 領域を表示
                    rectangle(frame, Rect(left, top, width, height), Scalar(0, 255, 0));

                    // ラベルとしてクラス番号と信頼度を表示
                    std::string label = format("%2d %.2f", classIdPoint.x, confidence);
                    int baseLine;
                    Size labelSize = getTextSize(label, FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);

                    top = max(top, labelSize.height);
                    rectangle(frame, Point(left, top - labelSize.height),
                        Point(left + labelSize.width, top + baseLine), Scalar::all(255), FILLED);
                    putText(frame, label, Point(left, top), FONT_HERSHEY_SIMPLEX, 0.5, Scalar());
                }
            }
        }
    }
}

카테고리명을 표시하는 부분은 이 기사에서는 남은 부분이므로 생략했습니다만, coco.names로부터 읽어 카테고리 번호로 옮겨놓으면 좋을 것입니다.