在C++ OpenCV中实现numpy.argmax

我正在使用OpenCV（4.6.0）DNN模块生成图像的语义分割，网络的输出是一个cv：：Mat，大小为（numberOfClasses x image_height x image_width），其中包含每个像素的类别概率。
我想计算每个像素的类ID的概率最高。
在Python中，函数numpy.argmax（src_matrix，axis=0）给出了所需的输出。
在C++ OpenCV中，函数cv：：reduceArgMax（src_，dst_，axis_）只能在二维矩阵上计算相同的东西。因此，我试图从三维矩阵中获取二维切片（（image_width x numberOfClasses）或（numberOfClasses x image_height）），并在循环中计算这些切片上的argmax。但是我无法获得正确的切片。
示例代码

int sizes[] = {numberOfClasses, imageHeight, imageWidth};
cv::Mat probabilityMatrix(3, sizes, CV_32F);
cv::Mat argMaxOfSlice(image_width, 1);

for(int r = 0; r < probabilityMatrix.size[1]; r++){

// cv::Mat Slice = get a 2D slice of the size (image_width x numberOfClasses) from the row r
// cv::reduceArgMax(Slice, argMaxOfSlice, 1);

...

}

最好，我只想使用OpenCV库，但我也可以使用Eigen（3.2.10）。

编辑：

Python示例代码沿着示例输入：

import numpy as np

# Shape of the example_input (3x3x4) where (ch, row, col)
example_input = np.array([[[ -1,  0,  -1,  0],
                           [ 0,  -1,  -1,  0],
                           [ 0,  -1,  -1,  -1]],

                          [[ -1,  -1,  1,  1],
                           [ -1,  -1,  -1,  -1],
                           [ 1,  -1,  1,  -1]],

                          [[ 2,  -1,  -1,  -1],
                           [ -1,  2,  2,  -1],
                           [ -1,  2,  -1,  2]]])

expected_output = np.array([[ 2,  0,  1,  1],
                            [ 0,  2,  2,  0],
                            [ 1,  2,  1,  2]])

function_output = np.argmax(example_input, axis=0)

if np.count_nonzero(expected_output - function_output) > 0 : 
    print("Something wrong")
else:
    print("Correct")

C++ OpenCV示例输入和预期输出

int example_size[3] = {3, 3, 4};
float example_input_data[36] = { -1,  0,  -1,  0, 0,  -1,  -1,  0,  0,  -1,  -1,  -1, -1,  -1,  1,  1, -1,  -1,  -1,  -1,
                            1,  -1,  1,  -1, 2,  -1,  -1,  -1, -1,  2,  2,  -1, -1,  2,  -1,  2};
        cv::Mat example_input (3, example_size, CV_32F,  example_input_data);

int expected_output_data[12] = { 2,  0,  1,  1, 0,  2,  2,  0, 1,  2,  1,  2};
cv::Mat expected_output (3, 4, CV_16U, expected_output_data);

谢谢你

感谢@DanMašek我做的实现如下：

cv::Mat reshaped = network_out.reshape(1, numberOfClasses);
cv::Mat argmax_row_matrix;
cv::reduceArgMax(reshaped, argmax_row_matrix, 0);
cv::Mat argmax_image_shape = argmax_row_matrix.reshape(1,rows);

但是，此实现的运行速度比以下实现慢：

cv::Mat classID = cv::Mat::zeros(rows, cols, CV_32S);
cv::Mat maxVal(rows, cols, CV_32F, network_out.data);

for (int ch = 0; ch < chns; ch++){
    for (int row = 0; row < rows; row++){
        const float *ptrScore = network_out.ptr<float>(0, ch, row);
        int *ptrMaxCl = classID.ptr<int>(row);
        float *ptrMaxVal = maxVal.ptr<float>(row);
        for (int col = 0; col < cols; col++){
            if (ptrScore[col] > ptrMaxVal[col]){
                ptrMaxVal[col] = ptrScore[col];
                ptrMaxCl[col] = ReduceClasses(ch);
            }
        }
    }
}

在C++ OpenCV中实现numpy.argmax

1条答案

相关问题

热门标签

最新问答