如何将DepthAI道路分割示例的Python Decode函数转为C++实现？

首先明确原Python逻辑的核心步骤：

• 移除batch维度（对应np.squeeze，将张量形状从(1, H, W, C)转为(H, W, C)）
• 对每个像素的所有通道取最大值对应的索引（对应np.argmax，得到H×W的类别索引图）
• 通过类别索引映射到对应颜色（对应np.take，生成彩色分割掩码）

以下是适配DepthAI C++ API的等价实现，假设你已获取模型输出的dai::Tensor对象，且颜色映射表为std::vector<cv::Vec3b> colors（每个元素对应类别的BGR颜色，适配OpenCV默认格式）：

cv::Mat decodeSegmentationMask(const dai::Tensor& outputTensor, const std::vector<cv::Vec3b>& colors) {
    // 获取张量形状：NHWC格式，即[batch, height, width, channels]
    auto shape = outputTensor.getShape();
    int height = shape[1];
    int width = shape[2];
    int numClasses = shape[3];

    // 获取张量数据，假设模型输出为FP16格式（DepthAI常见输出类型），转为float处理
    auto data = outputTensor.getData<float>();

    cv::Mat mask(height, width, CV_8UC3);

    // 遍历每个像素
    for (int y = 0; y < height; y++) {
        for (int x = 0; x < width; x++) {
            // 计算当前像素在张量中的起始索引（跳过batch维度）
            int startIdx = y * width * numClasses + x * numClasses;
            int maxIdx = 0;
            float maxVal = data[startIdx];

            // 找到当前像素通道中的最大值索引（实现argmax逻辑）
            for (int c = 1; c < numClasses; c++) {
                float val = data[startIdx + c];
                if (val > maxVal) {
                    maxVal = val;
                    maxIdx = c;
                }
            }

            // 映射颜色到当前像素（确保索引不越界）
            if (maxIdx < colors.size()) {
                mask.at<cv::Vec3b>(y, x) = colors[maxIdx];
            } else {
                // 索引超出范围时默认设为黑色
                mask.at<cv::Vec3b>(y, x) = cv::Vec3b(0, 0, 0);
            }
        }
    }

    return mask;
}

关键细节说明 #

• 张量数据格式适配：如果模型输出为INT8等其他类型，需将getData<float>()改为对应类型（如getData<int8_t>()），并同步调整数值比较逻辑。
• 颜色格式转换：若原Python代码用RGB颜色，需先转为OpenCV的BGR格式再存入cv::Vec3b，避免显示颜色颠倒。
• 性能优化：若需更高运行效率，可改用OpenCV矩阵批量操作或SIMD指令优化，上述代码以逻辑清晰为优先，满足基础场景需求。

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