Crnn.cpp 6.76 KB
Newer Older
Your Name's avatar
Your Name committed
1
2
3
4
5
6
7
8
9
10
#include <Crnn.h>
#include <migraphx/onnx.hpp>
#include <migraphx/gpu/target.hpp>
#include <Filesystem.h>
#include <SimpleLog.h>


namespace migraphxSamples
{

liucong's avatar
liucong committed
11
Crnn::Crnn()
Your Name's avatar
Your Name committed
12
13
14
15
16
17
18
19
{

}

Crnn::~Crnn()
{

    configurationFile.release();
shizhm's avatar
shizhm committed
20

Your Name's avatar
Your Name committed
21
22
23
24
}

ErrorCode Crnn::Initialize(InitializationParameterOfOcr initializationParameterOfOcr, bool dynamic)
{
liucong's avatar
liucong committed
25
26
27
    // 读取配置文件
    std::string configFilePath=initializationParameterOfOcr.configFilePath;
    if(Exists(configFilePath)==false)
Your Name's avatar
Your Name committed
28
    {
liucong's avatar
liucong committed
29
30
        LOG_ERROR(stdout, "no configuration file!\n");
        return CONFIG_FILE_NOT_EXIST;
Your Name's avatar
Your Name committed
31
    }
liucong's avatar
liucong committed
32
33
34
35
36
37
    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");
Your Name's avatar
Your Name committed
38
39

    // 获取配置文件参数
liucong's avatar
liucong committed
40
41
    cv::FileNode netNode = configurationFile["CrnnDynamic"];
    std::string modelPath=(std::string)netNode["ModelPath"];
Your Name's avatar
Your Name committed
42
43
44
45

    // 加载模型
    if(Exists(modelPath)==false)
    {
liucong's avatar
liucong committed
46
        LOG_ERROR(stdout,"%s not exist!\n",modelPath.c_str());
Your Name's avatar
Your Name committed
47
48
49
50
        return MODEL_NOT_EXIST;
    }

    if(dynamic)
shizhm's avatar
shizhm committed
51
    {
Your Name's avatar
Your Name committed
52
53
54
        migraphx::onnx_options onnx_options;
        onnx_options.map_input_dims["input"]={1,1,32,512};

shizhm's avatar
shizhm committed
55
        net = migraphx::parse_onnx(modelPath, onnx_options);
liucong's avatar
liucong committed
56
        LOG_INFO(stdout,"succeed to load model: %s\n",GetFileName(modelPath).c_str());
Your Name's avatar
Your Name committed
57

liucong's avatar
liucong committed
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
        // 获取模型输入/输出节点信息
        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.begin()->first;
        inputShape=inputs.begin()->second;
liucong's avatar
liucong committed
74
75
76
77
78
        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);
Your Name's avatar
Your Name committed
79
80

        // log输出日志信息
liucong's avatar
liucong committed
81
82
        LOG_INFO(stdout,"InputMaxSize:%dx%d\n",inputSize.width,inputSize.height);
        LOG_INFO(stdout,"InputName:%s\n",inputName.c_str());
Your Name's avatar
Your Name committed
83
84
85
86
87
88
    }
    else
    {
        migraphx::onnx_options onnx_options;
        onnx_options.map_input_dims["input"]={1,1,32,100};

shizhm's avatar
shizhm committed
89
        net = migraphx::parse_onnx(modelPath, onnx_options);
liucong's avatar
liucong committed
90
        LOG_INFO(stdout,"succeed to load model: %s\n",GetFileName(modelPath).c_str());
Your Name's avatar
Your Name committed
91

liucong's avatar
liucong committed
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
        // 获取模型输入/输出节点信息
        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.begin()->first;
        inputShape=inputs.begin()->second;
liucong's avatar
liucong committed
108
109
110
111
112
        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);
Your Name's avatar
Your Name committed
113
114

        // log输出日志信息
liucong's avatar
liucong committed
115
116
        LOG_INFO(stdout,"InputSize:%dx%d\n",inputSize.width,inputSize.height);
        LOG_INFO(stdout,"InputName:%s\n",inputName.c_str());
Your Name's avatar
Your Name committed
117
118
119
120
121
122
123
124
    }

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

    // 编译模型
    migraphx::compile_options options;
    options.device_id=0;                          // 设置GPU设备,默认为0号设备
shizhm's avatar
shizhm committed
125
126
    options.offload_copy=true;
    net.compile(gpuTarget,options);
liucong's avatar
liucong committed
127
    LOG_INFO(stdout,"succeed to compile model: %s\n",GetFileName(modelPath).c_str());
Your Name's avatar
Your Name committed
128

liucong's avatar
liucong committed
129
130
131
    // warm up
    std::unordered_map<std::string, migraphx::argument> inputData;
    inputData[inputName]=migraphx::argument{inputShape};
shizhm's avatar
shizhm committed
132
    net.eval(inputData);
Your Name's avatar
Your Name committed
133
134
135
136
137
138
139
140

    return SUCCESS;
}

ErrorCode Crnn::Infer(const cv::Mat &srcImage, std::vector<char> &resultsChar, bool raw, bool dynamic)
{
    if(srcImage.empty() || srcImage.type()!=CV_8UC3)
    {
liucong's avatar
liucong committed
141
        LOG_ERROR(stdout, "image error!\n");
Your Name's avatar
Your Name committed
142
143
144
145
146
        return IMAGE_ERROR;
    }

    cv::Mat inputImage, inputBlob;
    cv::cvtColor(srcImage, inputImage, CV_BGR2GRAY);
shizhm's avatar
shizhm committed
147

Your Name's avatar
Your Name committed
148
149
150
151
152
153
154
155
156
157
158
159
160
    int height, width, widthRaw;
    widthRaw = inputImage.cols;
    if(dynamic)
    {
        cv::resize(inputImage, inputImage, cv::Size(widthRaw, 32));
        height = inputImage.rows, width = inputImage.cols;
    }
    else
    {
        cv::resize(inputImage, inputImage, cv::Size(100, 32));
        height = inputImage.rows, width = inputImage.cols;
    }
    inputBlob = cv::dnn::blobFromImage(inputImage);
shizhm's avatar
shizhm committed
161

Your Name's avatar
Your Name committed
162
163
164
165
    for(int i=0; i<width * height; i++)
    {
        *((float*)inputBlob.data+i) = ((*((float*)inputBlob.data+i))/255.f - 0.5)/0.5;
    }
shizhm's avatar
shizhm committed
166

liucong's avatar
liucong committed
167
168
    // 创建输入数据
    std::unordered_map<std::string, migraphx::argument> inputData;
Your Name's avatar
Your Name committed
169
170
171
172
173
174
175
176
177
    if(dynamic)
    {
        std::vector<std::size_t> dynamicShape = {1, 1, 32, width};
        inputData[inputName]= migraphx::argument{migraphx::shape(inputShape.type(),dynamicShape), (float*)inputBlob.data};
    }
    else
    {
        inputData[inputName]= migraphx::argument{inputShape, (float*)inputBlob.data};
    }
shizhm's avatar
shizhm committed
178

Your Name's avatar
Your Name committed
179
180
    // 推理
    std::vector<migraphx::argument> inferenceResults = net.eval(inputData);
shizhm's avatar
shizhm committed
181

Your Name's avatar
Your Name committed
182
183
    // 获取推理结果
    std::vector<cv::Mat> outs;
shizhm's avatar
shizhm committed
184
    migraphx::argument result = inferenceResults[0];
Your Name's avatar
Your Name committed
185
186
187
188
189
190
191
192

    // 转换为cv::Mat
    migraphx::shape outputShape = result.get_shape();
    int shape[]={outputShape.lens()[0],outputShape.lens()[1],outputShape.lens()[2]};
    cv::Mat out(3,shape,CV_32F);
    memcpy(out.data,result.data(),sizeof(float)*outputShape.elements());
    outs.push_back(out);

shizhm's avatar
shizhm committed
193
    std::vector<int> predChars;
Your Name's avatar
Your Name committed
194
195
196
197
198
    const std::string alphabet = "-0123456789abcdefghijklmnopqrstuvwxyz";

    //获取字符索引序列
    for(uint i = 0; i < outs[0].size[0]; i++)
    {
liucong's avatar
liucong committed
199
        cv::Mat scores = cv::Mat(1,outs[0].size[2],CV_32F,outs[0].ptr<float>(i));
Your Name's avatar
Your Name committed
200
201
202
203
204
205
        cv::Point charIdPoint;
        double maxCharScore;
        cv::minMaxLoc(scores, 0, &maxCharScore, 0, &charIdPoint);
        int maxIdx = charIdPoint.x;
        predChars.push_back(maxIdx);
    }
shizhm's avatar
shizhm committed
206

Your Name's avatar
Your Name committed
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
    //字符转录处理
    for(uint i=0; i<predChars.size(); i++)
    {
        if(raw)
        {
            resultsChar.push_back(alphabet[predChars[i]]);
        }
        else
        {
            if(predChars[i] != 0)
            {
                if(!(i > 0 && predChars[i-1]==predChars[i]))
                {
                    resultsChar.push_back(alphabet[predChars[i]]);
                }
            }
        }
    }

    return SUCCESS;
}
}