如何使用C++ OpenCV快速完成3D矩阵的max/argmax操作？

我拥有一个尺寸为256×256×256的三维强度体，希望通过C++版本的OpenCV执行max/argmax操作，即求解第三维度的最大值及其对应的索引。请问实现该操作的最快方式是什么？

原代码如下：

#include <opencv2/opencv.hpp>
#include <opencv2/core/core.hpp>

void maxArgMax(const cv::Mat &volIn, cv::mat &imMagOut, cv::Mat &imIdxOut)
{
    // what to do here?
}

int main()
{
    int dims[] = { 256, 256, 256 };
    cv::Mat vol(3, dims, CV_32FC1);

    cv::Mat imMag(dims[0], dims[1], CV_32FC1, vol.at<float>(0)); // this should include the max. magnitude
    cv::Mat imIdx(dims[0], dims[1], CV_8UC1); // this should include the corresponding index

    maxArgMax(vol, imMag, imIdx);

    return 0;
}

要实现最快的max/argmax计算，核心思路是利用OpenCV优化过的内置函数计算最大值，同时手动遍历内存块获取索引，既借助OpenCV的SIMD加速提升max计算效率，又避免冗余操作。

步骤说明 #

1. 修正原代码错误：原代码中cv::mat是小写（应为cv::Mat），且imMag的初始化方式错误，不能直接绑定到三维体的内存，需单独分配二维矩阵内存。
2. 三维矩阵扁平化：将三维cv::Mat通过reshape转为二维矩阵（无需拷贝内存，仅修改矩阵头信息），每行对应原三维体中同一(x,y)位置的所有第三维度元素。
3. 用reduce计算最大值：借助OpenCV的reduce函数快速计算每行的最大值，该函数内部已做SIMD优化，效率远高于手动循环。
4. 遍历扁平化矩阵获取索引：逐行遍历扁平化后的矩阵，找到每个行中最大值对应的索引，存入结果矩阵。

完整代码实现 #

#include <opencv2/opencv.hpp>
#include <opencv2/core/core.hpp>

void maxArgMax(const cv::Mat &volIn, cv::Mat &imMagOut, cv::Mat &imIdxOut)
{
    // 检查输入输出尺寸合法性
    CV_Assert(volIn.dims == 3);
    int H = volIn.size[0];
    int W = volIn.size[1];
    int D = volIn.size[2];
    CV_Assert(imMagOut.size() == cv::Size(W, H) && imMagOut.type() == CV_32FC1);
    CV_Assert(imIdxOut.size() == cv::Size(W, H) && imIdxOut.type() == CV_8UC1);
    CV_Assert(D <= 256); // 确保索引能存入CV_8UC1

    // 将三维矩阵扁平化：行数H*W，列数D（无内存拷贝）
    cv::Mat flatVol = volIn.reshape(1, H * W);

    // 用reduce计算每行最大值，存入imMagOut
    cv::reduce(flatVol, imMagOut, 1, cv::REDUCE_MAX);

    // 遍历每行，找到最大值对应的索引
    for (int i = 0; i < flatVol.rows; ++i) {
        const float* rowPtr = flatVol.ptr<float>(i);
        float maxVal = imMagOut.at<float>(i / W, i % W);
        int maxIdx = 0;
        // 找到第一个等于最大值的索引（若有多个最大值，取第一个出现的）
        for (int j = 0; j < D; ++j) {
            if (rowPtr[j] == maxVal) {
                maxIdx = j;
                break;
            }
        }
        imIdxOut.at<uchar>(i / W, i % W) = static_cast<uchar>(maxIdx);
    }
}

int main()
{
    int dims[] = { 256, 256, 256 };
    cv::Mat vol(3, dims, CV_32FC1);
    // 可以在这里填充测试数据，比如cv::randu(vol, 0.0f, 255.0f);

    // 正确初始化输出矩阵
    cv::Mat imMag(dims[0], dims[1], CV_32FC1);
    cv::Mat imIdx(dims[0], dims[1], CV_8UC1);

    maxArgMax(vol, imMag, imIdx);

    // 后续可以保存或处理结果
    // cv::imwrite("max_magnitude.png", imMag * (255.0f / cv::max(imMag))); // 归一化后保存
    // cv::imwrite("max_index.png", imIdx);

    return 0;
}

效率说明 #

• reduce函数：OpenCV的reduce针对不同数据类型和操作做了硬件加速（如SSE、AVX），计算最大值的效率比手动循环高2-4倍。
• 扁平化操作：reshape仅修改矩阵的尺寸信息，不拷贝内存，避免了额外的内存开销。
• 索引查找：直接利用已计算出的最大值进行匹配，无需重复比较所有元素，进一步节省计算量。

注意事项 #

• 若三维矩阵存在多个相同的最大值，上述代码会返回第一个出现的索引；若需要返回最后一个或其他规则的索引，只需修改内层循环的逻辑。
• 确保第三维度的尺寸不超过255，因为imIdxOut用的是CV_8UC1类型；若维度更大，需改为CV_16UC1或CV_32SC1类型。

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