提交主要代码

3rdParty/InstallOpenCVDependences.sh

+#! /bin/sh
+
+############### Ubuntu  ###############
+# 参考：https://docs.opencv.org/3.4.11/d7/d9f/tutorial_linux_install.html
+apt-get install build-essential -y
+apt-get install cmake git libgtk2.0-dev pkg-config libavcodec-dev libavformat-dev libswscale-dev -y
+apt-get install python-dev python-numpy libtbb2 libtbb-dev libjpeg-dev libpng-dev libtiff-dev libjasper-dev libdc1394-22-dev -y # 处理图像所需的包，可选
+
+############### CentOS ###############
+# yum install gcc gcc-c++ gtk2-devel gimp-devel gimp-devel-tools gimp-help-browser zlib-devel libtiff-devel libjpeg-devel libpng-devel gstreamer-devel libavc1394-devel libraw1394-devel libdc1394-devel jasper-devel jasper-utils swig python libtool nasm -y
\ No newline at end of file

3rdParty/InstallRBuild.sh

+
+############################ 在线安装依赖 ###############################
+#cd ./3rdParty
+#pip install rbuild-master.tar.gz
+
+
+############################ 离线安装依赖 ###############################
+# 安装依赖
+cd ./3rdParty/rbuild_depend
+pip install click-6.6-py2.py3-none-any.whl
+pip install six-1.15.0-py2.py3-none-any.whl
+pip install subprocess32-3.5.4.tar.gz
+pip install cget-0.1.9.tar.gz
+
+# 安装rbuild
+cd ../
+pip install rbuild-master.tar.gz
+
+

Doc/Tutorial_Cpp.md

+# 图像分割
+
+本示例主要通过DeepLabv3模型说明如何使用MIGraphX C++ API进行图像分割模型的推理，包括模型初始化、预处理、模型推理。
+
+## 模型简介
+
+本示例采用了经典的DeepLabv3模型进行图像分割, 模型deeplabv3_resnet101.onnx文件保存在Resource/Models文件夹下。模型结构如下图所示，可以通过netron工具, 链接：https://netron.app/, 查看具体的模型结构，该模型的输入shape为[batch_size,3,513,513]，输出shape为[batch_size,513,513]，数据排布为NCHW。
+
+<img src="./Images/deeplabv3_01.jpg" style="zoom:80%;" align=middle>
+
+## 模型初始化
+
+在模型初始化的过程中，首先采用parse_onnx()函数根据提供的模型地址加载图像分割deeplabv3的onnx模型，保存在net中。其次，通过net.get_parameter_shapes()获取deeplabv3模型的输入属性，包含inputName和inputShape。最后，完成模型加载后使用migraphx::gpu::target{}设置编译模式为GPU模式，并使用compile()函数编译模型，完成模型的初始化过程。
+
+其中，模型地址设置在/Resource/Configuration.xml文件中的Unet节点中。
+
+```C++
+ErrorCode Unet::Initialize(InitializationParameterOfSegmentation initParamOfSegmentationUnet)
+{
+    ...
+
+    // 加载模型
+    net = migraphx::parse_onnx(modelPath); 
+    LOG_INFO(logFile,"succeed to load model: %s\n",GetFileName(modelPath).c_str());
+
+    // 获取模型输入/输出节点信息
+    std::cout<<"inputs:"<<std::endl;
+    std::unordered_map<std::string, migraphx::shape> inputs=net.get_inputs();
+    for(auto i:inputs)
+    {
+        std::cout<<i.first<<":"<<i.second<<std::endl;
+    }
+    std::cout<<"outputs:"<<std::endl;
+    std::unordered_map<std::string, migraphx::shape> outputs=net.get_outputs();
+    for(auto i:outputs)
+    {
+        std::cout<<i.first<<":"<<i.second<<std::endl;
+    }
+    inputName=inputs.first;
+    inputShape=inputs.second;
+    int N=inputShape.lens()[0];
+    int C=inputShape.lens()[1];
+    int H=inputShape.lens()[2];
+    int W=inputShape.lens()[3];
+    inputSize=cv::Size(W,H);
+
+    // 设置模型为GPU模式
+    migraphx::target gpuTarget = migraphx::gpu::target{};
+
+    // 编译模型
+    migraphx::compile_options options;
+    options.device_id=0;                          // 设置GPU设备，默认为0号设备
+    options.offload_copy=true;                    // 设置offload_copy
+    net.compile(gpuTarget,options);               
+    LOG_INFO(logFile,"succeed to compile model: %s\n",GetFileName(modelPath).c_str());
+
+    ...
+}
+```
+
+## 预处理
+
+完成模型初始化后，需要将输入数据进行如下预处理：
+
+​        1.尺度变换，将图像resize到513x513大小
+
+​        2.归一化，将数据归一化到[0.0, 1.0]之间
+
+​        3.数据排布，将数据从HWC转换为NCHW
+
+本示例代码主要通过opencv实现预处理操作：
+
+```C++
+ErrorCode Unet::Segmentation(const cv::Mat &srcImage, cv::Mat &maskImage)
+{
+    ...
+        
+    // 图像预处理并转换为NCHW
+    cv::Mat inputBlob;
+    cv::dnn::blobFromImage(srcImage,    // 输入数据，支持多张图像
+                    inputBlob,          // 输出数据
+                    1 / 255.0,          // 缩放系数
+                    inputSize,          // 模型输入大小，这里为256x256
+                    Scalar(0, 0, 0),    // 均值，这里不需要减均值，所以设置为0
+                    true,               // 通道转换，B通道与R通道互换，所以为true
+                    false);
+    
+    ...
+}
+```
+
+1.cv::dnn::blobFromImage()函数支持多个输入图像，首先将输入图像resize到inputSize大小，然后减去均值，其次乘以缩放系数1/255.0并转换为NCHW，最终将转换好的数据保存到inputBlob作为输入数据，执行后续的模型推理。
+
+## 推理
+
+完成图像预处理后，就可以执行模型推理。首先，定义inputData表示deeplabv3模型的输入数据，inputName表示deeplabv3模型的输入节点名，采用migraphx::argument{inputShape, (float*)inputBlob.data}保存前面预处理的数据inputBlob，第一个参数表示输入数据的shape，第二个参数表示输入数据指针。其次，执行net.eval(inputData)获得模型的推理结果，由于这里只有一个输出节点，仅使用results[0]获取输出节点的数据，就可以对输出数据执行相关后处理操作。
+
+```c++
+ErrorCode Unet::Segmentation(const cv::Mat &srcImage, cv::Mat &maskImage)
+{
+    ...
+        
+    // 创建输入数据
+    migraphx::parameter_map inputData;
+    inputData[inputName]= migraphx::argument{inputShape, (float*)inputBlob.data};
+
+    // 推理
+    std::vector<migraphx::argument> results = net.eval(inputData);
+    // 如果想要指定输出节点，可以给eval()函数中提供outputNames参数来实现
+    //std::vector<std::string> outputNames = {"outputs"};
+    //std::vector<migraphx::argument> inferenceResults = net.eval(inputData, outputNames);
+
+
+    // 获取输出节点的属性
+    migraphx::argument result = results[0];                 // 获取第一个输出节点的数据
+    migraphx::shape outputShape=result.get_shape();         // 输出节点的shape
+    std::vector<std::size_t> outputSize=outputShape.lens(); // 每一维大小，维度顺序为(N,C,H,W) 
+    int numberOfOutput=outputShape.elements();              // 输出节点元素的个数
+    float *data = (float *)result.data();                   // 输出节点数据指针
+
+    ...
+}
+```
+
+模型得到的推理结果并不能直接作为图像的分割结果，还需要做如下处理：
+
+1.计算softmax值，计算不同通道同一[H,W]位置的softmax值，找出概率最高的通道。
+
+2.保存结果，创建一个cv::Mat，根据不同的通道索引在颜色映射表取值并按行依次赋值到Mat对应位置，得到最终的分割图像。
+
+```c++
+ErrorCode Unet::Segmentation(const cv::Mat &srcImage, cv::Mat &maskImage)
+{
+    ...
+
+	cv::Mat outputImage(cv::Size(W, H), CV_8UC3);
+    // 创建颜色映射表
+    std::vector<cv::Scalar> color_map = create_color_map();
+
+    for(int i = 0;i < H; i++){
+        for(int j = 0;j < W;j++){
+            std::vector<float> channel_value;
+            for(int k = 0;k < C;k++){
+                channel_value.push_back(data[k*(H*W)+i*W+j]);
+            }
+            std::vector<float> probs = softmax(channel_value);
+            // 找到概率最高的类别索引
+            int max_index = std::max_element(probs.begin(),probs.end())-probs.begin();
+            cv::Scalar sc = color_map[max_index];
+            outputImage.at<cv::Vec3b>(i, j)[0]= sc.val[0]; 
+            outputImage.at<cv::Vec3b>(i, j)[1]= sc.val[1];
+            outputImage.at<cv::Vec3b>(i, j)[2]= sc.val[2];
+
+        }
+    }
+
+    maskImage = outputImage.clone();
+    
+    ...
+}
+```
+
+注：本次采用的模型权重onnx文件是通过使用PASCAL VOC 2012数据集来训练的。因此，“现实世界“图像的分割结果不完美是意料之中的。为了获得更好的结果，建议对现实世界示例数据集上的模型进行微调。
+

Doc/Tutorial_Python.md

+# 图像分割
+
+本示例主要通过deeplabv3模型说明如何使用MIGraphX Python API进行图像分割模型的推理，包括预处理、模型推理。
+
+## 模型简介
+
+本示例采用了经典的deeplabv3模型进行图像分割，将deeplabv3_resnet101.onnx文件保存在Resource/Models文件夹下。模型结构如下图所示, 链接：https://netron.app/, 该模型的输入shape为[batch_size,3,513,513]，输出shape为[batch_size,1,513,513]，数据排布为NCHW。
+
+<img src="./Images/deeplabv3_01.jpg" style="zoom:80%;" align=middle>
+
+## 预处理
+
+在将数据输入到模型之前，需要对图像做如下预处理操作：
+
+​      1.尺度变换，将图像resize到513x513大小
+
+​      2.数据排布，将数据从HWC转换为CHW，再将维度转换为NCHW
+
+​      3.归一化，将数据归一化到[0.0, 1.0]之间
+
+本示例代码通过如下方式实现预处理操作：
+
+```python
+def Preprocessing(pil_img, newW, newH):
+    assert newW > 0 and newH > 0, 'Scale is too small'  
+    img_nd = cv2.cvtColor(pil_img, cv2.COLOR_BGR2RGB)    # BGR转换为RGB
+    img_nd = cv2.resize(img_nd, (newW, newH))            # 将图像尺寸修改为256x256
+    
+    if len(img_nd.shape) == 2:
+        img_nd = np.expand_dims(img_nd, axis=2)
+
+    img_trans = img_nd.transpose((2, 0, 1))              # HWC转换为CHW
+    img_trans = np.expand_dims(img_trans, 0)             # CHW扩展为NCHW
+    img_trans = np.ascontiguousarray(img_trans)          # 保证内存连续存储
+    img_trans = img_trans.astype(np.float32)             # 转换成浮点型数据
+    if img_trans.max() > 1:                             
+        img = img_trans / 255.0                          # 保证数据处于0-1之间的浮点数
+    return img
+```
+
+## 推理
+
+完成图像预处理后，就可以执行推理，得到推理结果。
+
+```python
+    # 加载模型
+    model = migraphx.parse_onnx("../Resource/Models/deeplabv3_resnet101.onnx", map_input_dims=maxInput)
+
+    # 获取模型输入/输出节点信息
+    inputs = model.get_inputs()
+    outputs = model.get_outputs()
+
+    # 编译模型
+    model.compile(migraphx.get_target("gpu"), device_id=0)      # device_id: 设置GPU设备，默认为0号设备
+
+    # 图像预处理
+    img = cv2.imread("../Resource/Images/000001.jpg")
+    input_img = Preprocessing(img, 513, 513)
+
+    print(inputs)
+    # 模型推理
+    mask = model.run({'images':input_img})      
+    result = mask[0]                                            # 得到第一个输出节点的结果
+    # 对通道维度进行softmax
+    softmax_result = softmax(result)                            
+    # 计算通道维度最大值对应的索引（即类别索引）      
+    max_indices = np.argmax(softmax_result, axis=1)             # 等价于 np.argmax(arr, axis=1, keepdims=False)
+    # 使用预设颜色
+    color_map = np.array([
+        [0, 0, 0], [255, 0, 0], [0, 255, 0], [0, 0, 255],       # 0-3类
+        [255, 255, 0], [255, 0, 255], [0, 255, 255], [128, 0, 0],  # 4-7类
+        [0, 128, 0], [0, 0, 128], [128, 128, 0], [128, 0, 128],  # 8-11类
+        [0, 128, 128], [192, 192, 192], [128, 128, 128], [64, 0, 0],  # 12-15类
+        [0, 64, 0], [0, 0, 64], [64, 64, 0], [64, 0, 64],       # 16-19类
+        [0, 64, 64]                                             # 20类
+    ], dtype=np.uint8)
+
+    flat_index = max_indices[0]                                 # 取第0批的数据 
+    # 将二维的类别索引图直接转换为三维的 RGB 彩色图像
+    rgb_image = color_map[flat_index]
+    cv2.imwrite("Result.jpg", rgb_image)                         # 保存图像分割结果
+
+```
+
+1.Preprocessing函数返回预处理后的数据（numpy类型），然后通过model.run({'images':input_img})得到推理结果，因为只有输入一张图片，所以通过mask[0]获取第一个输出节点的数据即可。
+
+2.模型得到的推理结果并不能直接作为分割结果。首先，需要计算softmax值，计算不同通道同一[H,W]位置的softmax值，找出概率最高的通道。其次，根据不同的通道索引转换为三维的 RGB 彩色图像。最终，保存结果得到分割图像。
+
+注：本次采用的模型权重onnx文件是通过使用PASCAL VOC 2012数据集来训练的。因此，“现实世界“图像的分割结果不完美是意料之中的。为了获得更好的结果，建议对现实世界示例数据集上的模型进行微调。
+

Python/DeepLabV3.py

+import numpy as np
+import cv2
+import migraphx
+
+def Preprocessing(pil_img, newW, newH):
+    assert newW > 0 and newH > 0, 'Scale is too small' 
+    img_nd = cv2.cvtColor(pil_img, cv2.COLOR_BGR2RGB)    # BGR转换为RGB
+    img_nd = cv2.resize(img_nd, (newW, newH))            # 将图像尺寸修改为256x256
+    
+    if len(img_nd.shape) == 2:                           # 获取图像的维度信息
+        img_nd = np.expand_dims(img_nd, axis=2)          # 如果是2维的 扩充为3维  
+
+    img_trans = img_nd.transpose((2, 0, 1))              # HWC转换为CHW     
+    img_trans = np.expand_dims(img_trans, 0)             # CHW扩展为NCHW
+    img_trans = np.ascontiguousarray(img_trans)          # 保证内存连续存储
+    img_trans = img_trans.astype(np.float32)             # 转换成浮点型数据
+    if img_trans.max() > 1:                             
+        img = img_trans / 255.0                          # 保证数据处于0-1之间的浮点数
+
+    return img
+
+def Sigmoid(x):
+  return 1 / (1 + np.exp(-x))
+
+# 对通道维度执行Softmax
+def softmax(arr):
+    # 1：对通道维度计算指数，避免数值溢出（减去最大值）
+    exp_vals = np.exp(arr - np.max(arr, axis=1, keepdims=True))
+    # 2：计算通道维度的指数和
+    sum_exp = np.sum(exp_vals, axis=1, keepdims=True)
+    # 3：归一化得到Softmax结果
+    return exp_vals / sum_exp  
+
+if __name__ == '__main__':
+
+    # 设置最大输入shape
+    maxInput={"inputs":[1,3,256,256]}
+    
+    # 加载模型
+    model = migraphx.parse_onnx("../Resource/Models/deeplabv3_resnet101.onnx", map_input_dims=maxInput)
+
+    # 获取模型输入/输出节点信息
+    inputs = model.get_inputs()
+    outputs = model.get_outputs()
+
+    # 编译模型
+    model.compile(migraphx.get_target("gpu"), device_id=0)      # device_id: 设置GPU设备，默认为0号设备
+
+    # 图像预处理
+    img = cv2.imread("../Resource/Images/000001.jpg")
+    input_img = Preprocessing(img, 513, 513)
+
+    print(inputs)
+    # 模型推理
+    mask = model.run({'images':input_img})      
+    result = mask[0]                                            # 得到第一个输出节点的结果
+    # 对通道维度进行softmax
+    softmax_result = softmax(result)                            
+    # 计算通道维度最大值对应的索引（即类别索引）      
+    max_indices = np.argmax(softmax_result, axis=1)             # 等价于 np.argmax(arr, axis=1, keepdims=False)
+    # 使用预设颜色
+    color_map = np.array([
+        [0, 0, 0], [255, 0, 0], [0, 255, 0], [0, 0, 255],       # 0-3类
+        [255, 255, 0], [255, 0, 255], [0, 255, 255], [128, 0, 0],  # 4-7类
+        [0, 128, 0], [0, 0, 128], [128, 128, 0], [128, 0, 128],  # 8-11类
+        [0, 128, 128], [192, 192, 192], [128, 128, 128], [64, 0, 0],  # 12-15类
+        [0, 64, 0], [0, 0, 64], [64, 64, 0], [64, 0, 64],       # 16-19类
+        [0, 64, 64]                                             # 20类
+    ], dtype=np.uint8)
+
+    flat_index = max_indices[0]                                 # 取第0批的数据 
+    # 将二维的类别索引图直接转换为三维的 RGB 彩色图像
+    rgb_image = color_map[flat_index]
+    cv2.imwrite("Result.jpg", rgb_image)                         # 保存图像分割结果
+

Python/requirements.txt

+opencv-python
+numpy
\ No newline at end of file

Resource/Configuration.xml

+<?xml version="1.0" encoding="GB2312"?>
+<opencv_storage>
+
+	<!--DeepLabV3-->
+	<DeepLabV3>
+		<ModelPath>"../Resource/Models/deeplabv3_resnet101.onnx"</ModelPath>
+	</DeepLabV3>
+	
+</opencv_storage>

Src/DeepLabV3.cpp

+#include <DeepLabV3.h>
+
+#include <migraphx/onnx.hpp>
+#include <migraphx/gpu/target.hpp>
+
+#include <Filesystem.h>
+#include <SimpleLog.h>
+#include <algorithm>
+
+namespace migraphxSamples
+{
+
+static float Sigmoid(float x) { return (1 / (1 + exp(-x))); }
+
+static std::vector<float> softmax(vector<float> v){
+
+    std::vector<float> probs(v.size());
+    float max_val = *std::max_element(v.begin(),v.end());
+    float sum_exp = 0.0f;
+    for (int i = 0; i < v.size(); ++i) {
+        probs[i] = std::exp(v[i] - max_val);
+        sum_exp += probs[i];
+    }
+    for (int i = 0; i < v.size(); ++i) {
+       probs[i] /= sum_exp;
+    }
+
+    return probs;
+}
+
+// 定义21个类别的颜色映射表（BGR格式）
+std::vector<cv::Scalar> create_color_map() {
+    return {
+        cv::Scalar(0, 0, 0),       // 0: 黑色（背景）
+        cv::Scalar(255, 0, 0),     // 1: 蓝色
+        cv::Scalar(0, 255, 0),     // 2: 绿色
+        cv::Scalar(0, 0, 255),     // 3: 红色
+        cv::Scalar(255, 255, 0),   // 4: 青色
+        cv::Scalar(255, 0, 255),   // 5: 品红
+        cv::Scalar(0, 255, 255),   // 6: 黄色
+        cv::Scalar(128, 0, 0),     // 7: 深蓝
+        cv::Scalar(0, 128, 0),     // 8: 深绿
+        cv::Scalar(0, 0, 128),     // 9: 深红
+        cv::Scalar(128, 128, 0),   // 10: 深青
+        cv::Scalar(128, 0, 128),   // 11: 深品红
+        cv::Scalar(0, 128, 128),   // 12: 深黄
+        cv::Scalar(192, 192, 192), // 13: 灰色
+        cv::Scalar(128, 128, 128), // 14: 深灰
+        cv::Scalar(64, 0, 0),      // 15: 暗蓝
+        cv::Scalar(0, 64, 0),      // 16: 暗绿
+        cv::Scalar(0, 0, 64),      // 17: 暗红
+        cv::Scalar(64, 64, 0),     // 18: 暗青
+        cv::Scalar(64, 0, 64),     // 19: 暗品红
+        cv::Scalar(0, 64, 64)      // 20: 暗黄
+    };
+}
+
+
+
+DeepLabV3::DeepLabV3() {}
+
+DeepLabV3::~DeepLabV3() { configurationFile.release(); }
+
+
+ErrorCode DeepLabV3::Initialize(InitializationParameterOfSegmentation initParamOfSegmentationUnet)
+{
+    // 读取配置文件
+    std::string configFilePath = initParamOfSegmentationUnet.configFilePath;
+    if(!Exists(configFilePath))
+    {
+        LOG_ERROR(stdout, "no configuration file!\n");
+        return CONFIG_FILE_NOT_EXIST;
+    }
+    if(!configurationFile.open(configFilePath, cv::FileStorage::READ))
+    {
+        LOG_ERROR(stdout, "fail to open configuration file\n");
+        return FAIL_TO_OPEN_CONFIG_FILE;
+    }
+    LOG_INFO(stdout, "succeed to open configuration file\n");
+
+    // 获取配置文件参数
+    cv::FileNode netNode  = configurationFile["DeepLabV3"];
+    std::string modelPath = (std::string)netNode["ModelPath"];
+
+    // 设置最大输入shape
+    migraphx::onnx_options onnx_options;
+    onnx_options.map_input_dims["inputs"] = {1, 3, 256, 256};
+
+    // 加载模型
+    if(!Exists(modelPath))
+    {
+        LOG_ERROR(stdout, "%s not exist!\n", modelPath.c_str());
+        return MODEL_NOT_EXIST;
+    }
+    net = migraphx::parse_onnx(modelPath, onnx_options);
+    LOG_INFO(stdout, "succeed to load model: %s\n", GetFileName(modelPath).c_str());
+
+    // 获取模型输入/输出节点信息
+    std::unordered_map<std::string, migraphx::shape> inputs  = net.get_inputs();
+    std::unordered_map<std::string, migraphx::shape> outputs = net.get_outputs();
+    inputName = inputs.begin()->first;
+    inputShape = inputs.begin()->second;
+    int N = inputShape.lens()[0];
+    int C = inputShape.lens()[1];
+    int H = inputShape.lens()[2];
+    int W = inputShape.lens()[3];
+    inputSize = cv::Size(W, H);
+
+    // 设置模型为GPU模式
+    migraphx::target gpuTarget = migraphx::gpu::target{};
+
+    // 编译模型
+    migraphx::compile_options options;
+    options.device_id    = 0; // 设置GPU设备，默认为0号设备
+    options.offload_copy = true;
+    net.compile(gpuTarget, options);
+    LOG_INFO(stdout, "succeed to compile model: %s\n", GetFileName(modelPath).c_str());
+
+    // warm up
+    std::unordered_map<std::string, migraphx::argument> inputData;
+    inputData[inputName] = migraphx::argument{inputShape};
+    net.eval(inputData);
+
+    // log输出日志信息
+    LOG_INFO(stdout, "InputSize:%dx%d\n", inputSize.width, inputSize.height);
+    LOG_INFO(stdout, "InputName:%s\n", inputName.c_str());
+
+    return SUCCESS;
+}
+
+
+ErrorCode DeepLabV3::Segmentation(const cv::Mat& srcImage, cv::Mat& maskImage)
+{
+    if(srcImage.empty() || srcImage.type() != CV_8UC3)
+    {
+        LOG_ERROR(stdout, "image error!\n");
+        return IMAGE_ERROR;
+    }
+
+    // 数据预处理并转换为NCHW格式
+    cv::Mat inputBlob;
+    cv::dnn::blobFromImage(
+        srcImage, inputBlob, 1 / 255.0, inputSize, cv::Scalar(0, 0, 0), true, false);
+
+    // 创建输入数据
+    std::unordered_map<std::string, migraphx::argument> inputData;
+    inputData[inputName] = migraphx::argument{inputShape, (float*)inputBlob.data};
+
+    // 推理
+    std::vector<migraphx::argument> results = net.eval(inputData);
+
+    // 获取输出节点的属性
+    migraphx::argument result   = results[0];                 // 获取第一个输出节点的数据
+    migraphx::shape outputShape = result.get_shape();         // 输出节点的shape
+    std::vector<std::size_t> outputSize = outputShape.lens(); // 每一维大小，维度顺序为(N,C,H,W)
+
+    int numberOfOutput = outputShape.elements();              // 输出节点元素的个数
+    float* data        = (float*)result.data();               // 输出节点数据指针
+
+    int N = outputShape.lens()[0];
+    int C = outputShape.lens()[1];
+    int H = outputShape.lens()[2];
+    int W = outputShape.lens()[3];
+
+        
+        
+    cv::Mat outputImage(cv::Size(W, H), CV_8UC3);
+    // 创建颜色映射表
+    std::vector<cv::Scalar> color_map = create_color_map();
+
+    for(int i = 0;i < H; i++){
+        for(int j = 0;j < W;j++){
+            std::vector<float> channel_value;
+            for(int k = 0;k < C;k++){
+                channel_value.push_back(data[k*(H*W)+i*W+j]);
+            }
+            std::vector<float> probs = softmax(channel_value);
+            // 找到概率最高的类别索引
+            int max_index = std::max_element(probs.begin(),probs.end())-probs.begin();
+            cv::Scalar sc = color_map[max_index];
+            outputImage.at<cv::Vec3b>(i, j)[0]= sc.val[0]; 
+            outputImage.at<cv::Vec3b>(i, j)[1]= sc.val[1];
+            outputImage.at<cv::Vec3b>(i, j)[2]= sc.val[2];
+
+        }
+    }
+
+    maskImage = outputImage.clone();
+    return SUCCESS;
+}
+
+}
\ No newline at end of file