在C++ OpenCV中实现类似numpy.argmax(axis=0)的语义分割处理

问题背景 #

使用OpenCV 4.6.0的DNN模块进行图像语义分割时，网络输出为**(类别数 × 图像高度 × 图像宽度)**的3D cv::Mat，存储每个像素对应各类别的概率。需要为每个像素找出概率最高的类别ID，等价于Python中numpy.argmax(example_input, axis=0)的功能，但OpenCV原生的cv::reduceArgMax仅支持2D矩阵，手动提取2D切片时出现问题。优先采用OpenCV原生方法，也可使用Eigen 3.2.10实现。

示例代码 #

Python参考实现

import numpy as np

# 输入形状为(类别数, 高度, 宽度) = (3, 3, 4)
example_input = np.array([[[ -1,  0,  -1,  -1],
                           [ 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]]])

# 预期输出：(高度, 宽度) = (3,4)，每个值是对应像素的最大概率类别ID
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("结果错误")
else:
    print("结果正确")

C++ OpenCV输入与预期输出

#include <opencv2/opencv.hpp>

int main() {
    // 输入形状：(3类别, 3行, 4列)
    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);

    // 预期输出：(3行,4列)的类别ID矩阵
    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);

    // 解决方案代码写在这里...
    return 0;
}

解决方案1：OpenCV原生实现（推荐） #

核心思路：将3D矩阵**重塑(reshape)**为2D矩阵，利用cv::reduce的REDUCE_MAX_IDX模式计算每列的最大值索引，最后再重塑回目标形状。

实现代码 #

在上述C++示例的main函数中添加以下代码：

// 1. 获取3D矩阵的维度信息
int num_classes = example_input.size[0];
int height = example_input.size[1];
int width = example_input.size[2];

// 2. 将3D矩阵reshape为2D：(num_classes, height*width)
cv::Mat flat_input = example_input.reshape(1, num_classes);

// 3. 对每一列（对应一个像素的所有类别概率）求最大值索引
cv::Mat flat_output;
cv::reduce(flat_input, flat_output, 0, cv::REDUCE_MAX_IDX, CV_16U);

// 4. 将结果reshape为(height, width)的目标矩阵
cv::Mat function_output = flat_output.reshape(1, height);

// 5. 验证结果
cv::Mat diff;
cv::absdiff(function_output, expected_output, diff);
if (cv::countNonZero(diff) == 0) {
    std::cout << "结果正确" << std::endl;
} else {
    std::cout << "结果错误" << std::endl;
}

原理说明 #

• OpenCV的3D矩阵在内存中是连续存储的，reshape操作不会改变内存布局，只是修改矩阵的维度描述，效率极高。
• cv::reduce的REDUCE_MAX_IDX模式会计算指定轴上的最大值对应的索引，正好对应我们需要的类别ID。

解决方案2：Eigen实现 #

如果需要使用Eigen，可将cv::Mat转换为Eigen张量，利用Eigen的argmax函数实现。

实现代码 #

#include <opencv2/opencv.hpp>
#include <Eigen/Dense>
#include <Eigen/Tensor>

int main() {
    // （重复输入定义代码...）

    int num_classes = example_input.size[0];
    int height = example_input.size[1];
    int width = example_input.size[2];

    // 1. 将cv::Mat转换为Eigen张量
    Eigen::TensorMap<Eigen::Tensor<float, 3>> eigen_input(
        reinterpret_cast<float*>(example_input.data()),
        num_classes, height, width
    );

    // 2. 在类别维度（第0轴）上求argmax，得到(height, width)的结果
    Eigen::Tensor<int, 2> eigen_output = eigen_input.argmax(0);

    // 3. 将Eigen张量转换为cv::Mat
    cv::Mat function_output(height, width, CV_16U);
    memcpy(function_output.data(), eigen_output.data(), eigen_output.size() * sizeof(int));

    // 4. 验证结果
    cv::Mat diff;
    cv::absdiff(function_output, expected_output, diff);
    if (cv::countNonZero(diff) == 0) {
        std::cout << "结果正确" << std::endl;
    } else {
        std::cout << "结果错误" << std::endl;
    }

    return 0;
}

注意事项 #

• 需要确保Eigen版本为3.2.10及以上（该版本已支持张量的argmax操作）。
• 内存拷贝操作需保证数据类型匹配，这里类别ID用CV_16U存储，和Eigen的int兼容。

内容的提问来源于stack exchange，提问作者oguzhan ilter