Opencv使用Dnn模块推理Deeplab/Unet等语义分割模型

Python

import cv2
import numpy as np
img = cv2.imread("bus.jpg")
model_file = "deeplab.onnx"
if img is None:
    print("can not find image！")
org_height, org_width = img.shape[:2]
scale = 512 / max(org_width, org_height)
resized_img = cv2.resize(img, (0, 0), fx=scale, fy=scale)
resized_img = cv2.cvtColor(resized_img, cv2.COLOR_BGR2RGB)
pad_x = (512 - resized_img.shape[1]) // 2
pad_y = (512 - resized_img.shape[0]) // 2
padded_img = cv2.copyMakeBorder(
        resized_img, pad_y, pad_y, pad_x, pad_x,
        cv2.BORDER_CONSTANT, value=(128, 128, 128))
padded_img = padded_img.astype(np.float32) / 255.0
padded_img = cv2.resize(padded_img, (512, 512), interpolation=cv2.INTER_LINEAR)
input_blob = cv2.dnn.blobFromImage(padded_img)
net = cv2.dnn.readNet(model_file)
net.setInput(input_blob)
output = net.forward(net.getUnconnectedOutLayersNames())
pdata = output[0].squeeze()
num_cn, out_h, out_w = pdata.shape
color_table = [np.array([0, 0, 0], dtype=np.uint8)]
rng = np.random.default_rng()
fpdata = output[0].squeeze()
num_cn, out_h, out_w = pdata.shape
color_table = [np.array([0, 0, 0], dtype=np.uint8)]
rng = np.random.default_rng()
for i in range(1, num_cn):
    color_table.append(rng.integers(0, 256, size=3, dtype=np.uint8))
result = np.zeros((out_h, out_w, 3), dtype=np.uint8)
for row in range(out_h):
    for col in range(out_w):
        max_index = np.argmax(pdata[:, row, col])
        result[row, col] = color_table[max_index]
mask = result[pad_y:pad_y + int(org_height * scale), pad_x:pad_x + int(org_width * scale)]
mask = cv2.resize(mask, (org_width, org_height))
cv2.addWeighted(img, 0.7, mask, 0.3, 0, img)
cv2.imshow("result", img)
cv2.waitKey(0)

Cpp

#include <opencv2/dnn.hpp>
#include <opencv2/opencv.hpp>
#include <iostream>
using namespace cv;
using namespace std;
int main() {
    Mat img = imread("bus.jpg");
    string modelFile = "deeplab.onnx";
    if (img.empty()) {
        cout << "can not find image！" << endl;
        return -1;
    }
    int org_width = img.cols;
    int org_height = img.rows;
    double scale = 1.0;
    if (org_width >= org_height) {
        scale = (double)512 / org_width;
    }
    else {
        scale = (double)512 / org_height;
    }
    Mat resized_img;
    resize(img, resized_img, Size(), scale, scale);
    cv::cvtColor(resized_img, resized_img, cv::COLOR_BGR2RGB);
    int pad_x = (512 - resized_img.cols) / 2;
    int pad_y = (512 - resized_img.rows) / 2;
    Mat padded_img = Mat::zeros(512, 512, CV_8UC3);
    cv::cvtColor(img, padded_img, cv::COLOR_BGR2RGB);
    copyMakeBorder(resized_img, padded_img, pad_y, pad_y, pad_x, pad_x, BORDER_CONSTANT, Scalar(128,128,128));
    padded_img.convertTo(padded_img, CV_32FC3, 1.f / 255.f);
    cv::resize(padded_img, padded_img, cv::Size(512, 512), 0, 0, cv::INTER_LINEAR);
    Mat input = dnn::blobFromImage(padded_img);
    std::vector<cv::Mat> output;
    dnn::Net net = cv::dnn::readNet(modelFile);
    net.setInput(input);
    net.forward(output, net.getUnconnectedOutLayersNames());
    float* pdata = (float*)output[0].data;
    int num_cn = output[0].size[1];
    int out_h = output[0].size[2];
    int out_w = output[0].size[3];
    std::vector<cv::Vec3b> color_table;
    color_table.push_back(cv::Vec3b(0, 0, 0));
    cv::RNG rng;
    std::vector<uchar> color_tables;
    for (int i = 1; i < num_cn; i++) {
        color_table.push_back(cv::Vec3b(rng.uniform(0, 255), rng.uniform(0, 255), rng.uniform(0, 255)));
    }
    cv::Mat result = cv::Mat::zeros(cv::Size(out_w, out_h), CV_8UC3);
    for (int row = 0; row < out_h; row++) {
        for (int col = 0; col < out_w; col++) {
            int max_index = 0;
            float max_prob = pdata[row * out_w + col];
            for (int cn = 1; cn < num_cn; cn++) {
                float prob = pdata[cn * out_h * out_w + row * out_w + col];
                if (prob > max_prob) {
                    max_prob = prob;
                    max_index = cn;
                }
            }
            result.at<cv::Vec3b>(row, col) = color_table[max_index];
        }
    }
    Mat mask = result(Range(pad_y, pad_y + org_height * scale), Range(pad_x, pad_x + org_width * scale));
    resize(mask, mask, Size(org_width, org_height));
    cv::addWeighted(img, 0.7, mask, 0.3, 0, img);
    cv::imshow("result", img);
    cv::waitKey(0);
    return 0;
}

CMakeLists

cmake_minimum_required(VERSION 3.18)
project(Seg)
set("OpenCV_DIR" "E:\\Opencv\\opencv_vs\\build")
set(OpenCV_INCLUDE_DIRS ${OpenCV_DIR}\\include)
set(OpenCV_LIB_DIRS ${OpenCV_DIR}\\x64\\vc16\\lib) 
set(OpenCV_LIBS "opencv_world480d.lib" "opencv_world480.lib")    
set(CMAKE_CXX_STANDARD 14)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
include_directories(${OpenCV_INCLUDE_DIRS}) 
link_directories(${OpenCV_LIB_DIRS})  
find_package(OpenCV QUIET)	
link_libraries(${OpenCV_LIBS})
add_executable(Seg main.cpp)
target_compile_features(Seg PRIVATE cxx_std_14)