Unet.cpp 4.85 KB
Newer Older
liucong's avatar
liucong committed
1
2
3
4
5
6
7
8
9
10
11
#include <Unet.h>

#include <migraphx/onnx.hpp>
#include <migraphx/gpu/target.hpp>

#include <Filesystem.h>
#include <SimpleLog.h>

namespace migraphxSamples
{

liucong's avatar
liucong committed
12
static float Sigmoid(float x) { return (1 / (1 + exp(-x))); }
liucong's avatar
liucong committed
13

liucong's avatar
liucong committed
14
Unet::Unet() {}
liucong's avatar
liucong committed
15

liucong's avatar
liucong committed
16
Unet::~Unet() { configurationFile.release(); }
liucong's avatar
liucong committed
17
18
19
20

ErrorCode Unet::Initialize(InitializationParameterOfSegmentation initParamOfSegmentationUnet)
{
    // 读取配置文件
liucong's avatar
liucong committed
21
    std::string configFilePath = initParamOfSegmentationUnet.configFilePath;
liucong's avatar
liucong committed
22
    if(!Exists(configFilePath))
liucong's avatar
liucong committed
23
24
25
26
27
28
    {
        LOG_ERROR(stdout, "no configuration file!\n");
        return CONFIG_FILE_NOT_EXIST;
    }
    if(!configurationFile.open(configFilePath, cv::FileStorage::READ))
    {
liucong's avatar
liucong committed
29
30
        LOG_ERROR(stdout, "fail to open configuration file\n");
        return FAIL_TO_OPEN_CONFIG_FILE;
liucong's avatar
liucong committed
31
32
33
34
    }
    LOG_INFO(stdout, "succeed to open configuration file\n");

    // 获取配置文件参数
liucong's avatar
liucong committed
35
36
    cv::FileNode netNode  = configurationFile["Unet"];
    std::string modelPath = (std::string)netNode["ModelPath"];
liucong's avatar
liucong committed
37

liucong's avatar
liucong committed
38
39
    // 设置最大输入shape
    migraphx::onnx_options onnx_options;
liucong's avatar
liucong committed
40
    onnx_options.map_input_dims["inputs"] = {1, 3, 256, 256};
liucong's avatar
liucong committed
41

liucong's avatar
liucong committed
42
    // 加载模型
liucong's avatar
liucong committed
43
    if(!Exists(modelPath))
liucong's avatar
liucong committed
44
    {
liucong's avatar
liucong committed
45
        LOG_ERROR(stdout, "%s not exist!\n", modelPath.c_str());
liucong's avatar
liucong committed
46
47
        return MODEL_NOT_EXIST;
    }
liucong's avatar
liucong committed
48
49
    net = migraphx::parse_onnx(modelPath, onnx_options);
    LOG_INFO(stdout, "succeed to load model: %s\n", GetFileName(modelPath).c_str());
liucong's avatar
liucong committed
50

liucong's avatar
liucong committed
51
    // 获取模型输入/输出节点信息
liucong's avatar
liucong committed
52
53
54
55
56
57
58
59
60
    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);
liucong's avatar
liucong committed
61
62
63
64
65
66

    // 设置模型为GPU模式
    migraphx::target gpuTarget = migraphx::gpu::target{};

    // 编译模型
    migraphx::compile_options options;
liucong's avatar
liucong committed
67
68
69
70
    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());
liucong's avatar
liucong committed
71
72
73

    // warm up
    std::unordered_map<std::string, migraphx::argument> inputData;
liucong's avatar
liucong committed
74
75
    inputData[inputName] = migraphx::argument{inputShape};
    net.eval(inputData);
liucong's avatar
liucong committed
76
77

    // log输出日志信息
liucong's avatar
liucong committed
78
79
    LOG_INFO(stdout, "InputSize:%dx%d\n", inputSize.width, inputSize.height);
    LOG_INFO(stdout, "InputName:%s\n", inputName.c_str());
liucong's avatar
liucong committed
80
81
82
83

    return SUCCESS;
}

liucong's avatar
liucong committed
84
ErrorCode Unet::Segmentation(const cv::Mat& srcImage, cv::Mat& maskImage)
liucong's avatar
liucong committed
85
{
liucong's avatar
liucong committed
86
    if(srcImage.empty() || srcImage.type() != CV_8UC3)
liucong's avatar
liucong committed
87
88
89
90
91
92
93
    {
        LOG_ERROR(stdout, "image error!\n");
        return IMAGE_ERROR;
    }

    // 数据预处理并转换为NCHW格式
    cv::Mat inputBlob;
liucong's avatar
liucong committed
94
95
    cv::dnn::blobFromImage(
        srcImage, inputBlob, 1 / 255.0, inputSize, cv::Scalar(0, 0, 0), true, false);
liucong's avatar
liucong committed
96
97
98

    // 创建输入数据
    std::unordered_map<std::string, migraphx::argument> inputData;
liucong's avatar
liucong committed
99
    inputData[inputName] = migraphx::argument{inputShape, (float*)inputBlob.data};
liucong's avatar
liucong committed
100
101
102
103
104

    // 推理
    std::vector<migraphx::argument> results = net.eval(inputData);

    // 获取输出节点的属性
liucong's avatar
liucong committed
105
106
107
108
109
    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();               // 输出节点数据指针
liucong's avatar
liucong committed
110
111
112

    // 计算sigmoid值
    int value_mask[numberOfOutput];
liucong's avatar
liucong committed
113
    for(int i = 0; i < numberOfOutput; ++i)
liucong's avatar
liucong committed
114
    {
liucong's avatar
liucong committed
115
116
        float num = Sigmoid(data[i]);
        if(num > 0.996)
liucong's avatar
liucong committed
117
118
119
120
121
122
123
124
125
126
        {
            value_mask[i] = 1;
        }
        else
        {
            value_mask[i] = 0;
        }
    }

    // 将对应的value_mask[]数组中的值按行依次赋值到outputImage对应位置处
liucong's avatar
liucong committed
127
128
129
    cv::Mat outputImage =
        cv::Mat_<int>(cv::Size(outputShape.lens()[3], outputShape.lens()[2]), CV_32S);
    for(int i = 0; i < outputShape.lens()[2]; ++i)
liucong's avatar
liucong committed
130
    {
liucong's avatar
liucong committed
131
        for(int j = 0; j < outputShape.lens()[3]; ++j)
liucong's avatar
liucong committed
132
        {
liucong's avatar
liucong committed
133
            outputImage.at<int>(i, j) = value_mask[256 * i + j];
liucong's avatar
liucong committed
134
135
136
137
138
139
140
        }
    }
    outputImage.convertTo(maskImage, CV_8U, 255.0);

    return SUCCESS;
}

liucong's avatar
liucong committed
141
} // namespace migraphxSamples