paddlex.cpp 30.5 KB
Newer Older
C
Channingss 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13
//   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
S
syyxsxx 已提交
14 15

#include <math.h>
J
jack 已提交
16
#include <omp.h>
J
jack 已提交
17
#include <algorithm>
J
jack 已提交
18
#include <fstream>
J
jack 已提交
19
#include <cstring>
J
jack 已提交
20
#include "include/paddlex/paddlex.h"
S
syyxsxx 已提交
21 22 23 24 25

#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>

C
Channingss 已提交
26 27 28 29
namespace PaddleX {

void Model::create_predictor(const std::string& model_dir,
                             bool use_gpu,
C
Channingss 已提交
30
                             bool use_trt,
S
syyxsxx 已提交
31
                             bool use_mkl,
C
Channingss 已提交
32
                             int gpu_id,
J
jack 已提交
33
                             std::string key,
S
syyxsxx 已提交
34
                             int mkl_thread_num,
J
jack 已提交
35
                             bool use_ir_optim) {
C
Channingss 已提交
36 37 38
  paddle::AnalysisConfig config;
  std::string model_file = model_dir + OS_PATH_SEP + "__model__";
  std::string params_file = model_dir + OS_PATH_SEP + "__params__";
J
jack 已提交
39
  std::string yaml_file = model_dir + OS_PATH_SEP + "model.yml";
J
jack 已提交
40
  std::string yaml_input = "";
C
Channingss 已提交
41
#ifdef WITH_ENCRYPTION
J
jack 已提交
42
  if (key != "") {
F
FlyingQianMM 已提交
43 44
    model_file = model_dir + OS_PATH_SEP + "__model__.encrypted";
    params_file = model_dir + OS_PATH_SEP + "__params__.encrypted";
J
jack 已提交
45
    yaml_file = model_dir + OS_PATH_SEP + "model.yml.encrypted";
J
jack 已提交
46 47
    paddle_security_load_model(
        &config, key.c_str(), model_file.c_str(), params_file.c_str());
J
jack 已提交
48
    yaml_input = decrypt_file(yaml_file.c_str(), key.c_str());
C
Channingss 已提交
49 50
  }
#endif
J
jack 已提交
51
  if (yaml_input == "") {
S
syyxsxx 已提交
52
    // read yaml file
J
jack 已提交
53 54 55 56 57 58 59
    std::ifstream yaml_fin(yaml_file);
    yaml_fin.seekg(0, std::ios::end);
    size_t yaml_file_size = yaml_fin.tellg();
    yaml_input.assign(yaml_file_size, ' ');
    yaml_fin.seekg(0);
    yaml_fin.read(&yaml_input[0], yaml_file_size);
  }
S
syyxsxx 已提交
60
  // load yaml file
J
jack 已提交
61
  if (!load_config(yaml_input)) {
J
jack 已提交
62 63 64 65
    std::cerr << "Parse file 'model.yml' failed!" << std::endl;
    exit(-1);
  }

J
jack 已提交
66
  if (key == "") {
C
Channingss 已提交
67 68
    config.SetModel(model_file, params_file);
  }
S
syyxsxx 已提交
69 70
  if (use_mkl && name != "HRNet" && name != "DeepLabv3p") {
    config.EnableMKLDNN();
S
syyxsxx 已提交
71
    config.SetCpuMathLibraryNumThreads(mkl_thread_num);
S
syyxsxx 已提交
72
  }
C
Channingss 已提交
73 74 75 76 77 78 79
  if (use_gpu) {
    config.EnableUseGpu(100, gpu_id);
  } else {
    config.DisableGpu();
  }
  config.SwitchUseFeedFetchOps(false);
  config.SwitchSpecifyInputNames(true);
S
syyxsxx 已提交
80
  // enable graph Optim
F
FlyingQianMM 已提交
81 82 83
#if defined(__arm__) || defined(__aarch64__)
  config.SwitchIrOptim(false);
#else
J
jack 已提交
84
  config.SwitchIrOptim(use_ir_optim);
F
FlyingQianMM 已提交
85
#endif
S
syyxsxx 已提交
86
  // enable Memory Optim
C
Channingss 已提交
87
  config.EnableMemoryOptim();
C
Channingss 已提交
88 89 90 91 92 93 94 95
  if (use_trt) {
    config.EnableTensorRtEngine(
        1 << 20 /* workspace_size*/,
        32 /* max_batch_size*/,
        20 /* min_subgraph_size*/,
        paddle::AnalysisConfig::Precision::kFloat32 /* precision*/,
        true /* use_static*/,
        false /* use_calib_mode*/);
C
Channingss 已提交
96
  }
C
Channingss 已提交
97 98 99
  predictor_ = std::move(CreatePaddlePredictor(config));
}

J
jack 已提交
100 101
bool Model::load_config(const std::string& yaml_input) {
  YAML::Node config = YAML::Load(yaml_input);
C
Channingss 已提交
102 103
  type = config["_Attributes"]["model_type"].as<std::string>();
  name = config["Model"].as<std::string>();
F
FlyingQianMM 已提交
104 105
  std::string version = config["version"].as<std::string>();
  if (version[0] == '0') {
J
jack 已提交
106 107 108 109 110
    std::cerr << "[Init] Version of the loaded model is lower than 1.0.0, "
              << "deployment cannot be done, please refer to "
              << "https://github.com/PaddlePaddle/PaddleX/blob/develop/docs"
              << "/tutorials/deploy/upgrade_version.md "
              << "to transfer version." << std::endl;
F
FlyingQianMM 已提交
111 112
    return false;
  }
C
Channingss 已提交
113 114 115 116 117 118 119 120 121 122 123
  bool to_rgb = true;
  if (config["TransformsMode"].IsDefined()) {
    std::string mode = config["TransformsMode"].as<std::string>();
    if (mode == "BGR") {
      to_rgb = false;
    } else if (mode != "RGB") {
      std::cerr << "[Init] Only 'RGB' or 'BGR' is supported for TransformsMode"
                << std::endl;
      return false;
    }
  }
S
syyxsxx 已提交
124
  // build data preprocess stream
C
Channingss 已提交
125
  transforms_.Init(config["Transforms"], to_rgb);
S
syyxsxx 已提交
126
  // read label list
C
Channingss 已提交
127 128 129 130 131 132 133 134 135 136
  labels.clear();
  for (const auto& item : config["_Attributes"]["labels"]) {
    int index = labels.size();
    labels[index] = item.as<std::string>();
  }
  return true;
}

bool Model::preprocess(const cv::Mat& input_im, ImageBlob* blob) {
  cv::Mat im = input_im.clone();
137
  if (!transforms_.Run(&im, blob)) {
C
Channingss 已提交
138 139 140 141 142
    return false;
  }
  return true;
}

J
jack 已提交
143
// use openmp
J
jack 已提交
144 145 146
bool Model::preprocess(const std::vector<cv::Mat>& input_im_batch,
                       std::vector<ImageBlob>* blob_batch,
                       int thread_num) {
J
jack 已提交
147
  int batch_size = input_im_batch.size();
J
jack 已提交
148
  bool success = true;
J
jack 已提交
149 150
  thread_num = std::min(thread_num, batch_size);
  #pragma omp parallel for num_threads(thread_num)
J
jack 已提交
151
  for (int i = 0; i < input_im_batch.size(); ++i) {
J
jack 已提交
152
    cv::Mat im = input_im_batch[i].clone();
J
jack 已提交
153
    if (!transforms_.Run(&im, &(*blob_batch)[i])) {
J
jack 已提交
154 155 156 157 158 159
      success = false;
    }
  }
  return success;
}

C
Channingss 已提交
160 161 162 163 164
bool Model::predict(const cv::Mat& im, ClsResult* result) {
  inputs_.clear();
  if (type == "detector") {
    std::cerr << "Loading model is a 'detector', DetResult should be passed to "
                 "function predict()!"
J
jack 已提交
165
                 "to function predict()!" << std::endl;
C
Channingss 已提交
166 167
    return false;
  }
S
syyxsxx 已提交
168
  // im preprocess
C
Channingss 已提交
169 170 171 172
  if (!preprocess(im, &inputs_)) {
    std::cerr << "Preprocess failed!" << std::endl;
    return false;
  }
S
syyxsxx 已提交
173
  // predict
C
Channingss 已提交
174 175 176 177 178 179
  auto in_tensor = predictor_->GetInputTensor("image");
  int h = inputs_.new_im_size_[0];
  int w = inputs_.new_im_size_[1];
  in_tensor->Reshape({1, 3, h, w});
  in_tensor->copy_from_cpu(inputs_.im_data_.data());
  predictor_->ZeroCopyRun();
S
syyxsxx 已提交
180
  // get result
C
Channingss 已提交
181 182 183 184 185 186 187 188 189
  auto output_names = predictor_->GetOutputNames();
  auto output_tensor = predictor_->GetOutputTensor(output_names[0]);
  std::vector<int> output_shape = output_tensor->shape();
  int size = 1;
  for (const auto& i : output_shape) {
    size *= i;
  }
  outputs_.resize(size);
  output_tensor->copy_to_cpu(outputs_.data());
S
syyxsxx 已提交
190
  // postprocess
C
Channingss 已提交
191 192 193 194
  auto ptr = std::max_element(std::begin(outputs_), std::end(outputs_));
  result->category_id = std::distance(std::begin(outputs_), ptr);
  result->score = *ptr;
  result->category = labels[result->category_id];
J
jack 已提交
195
  return true;
C
Channingss 已提交
196 197
}

J
jack 已提交
198 199 200 201
bool Model::predict(const std::vector<cv::Mat>& im_batch,
                    std::vector<ClsResult>* results,
                    int thread_num) {
  for (auto& inputs : inputs_batch_) {
J
jack 已提交
202 203 204 205
    inputs.clear();
  }
  if (type == "detector") {
    std::cerr << "Loading model is a 'detector', DetResult should be passed to "
J
jack 已提交
206
                 "function predict()!" << std::endl;
J
jack 已提交
207 208 209
    return false;
  } else if (type == "segmenter") {
    std::cerr << "Loading model is a 'segmenter', SegResult should be passed "
J
jack 已提交
210
                 "to function predict()!" << std::endl;
J
jack 已提交
211 212
    return false;
  }
J
jack 已提交
213
  inputs_batch_.assign(im_batch.size(), ImageBlob());
S
syyxsxx 已提交
214
  // preprocess
J
jack 已提交
215
  if (!preprocess(im_batch, &inputs_batch_, thread_num)) {
J
jack 已提交
216 217 218
    std::cerr << "Preprocess failed!" << std::endl;
    return false;
  }
S
syyxsxx 已提交
219
  // predict
J
jack 已提交
220 221 222 223 224 225
  int batch_size = im_batch.size();
  auto in_tensor = predictor_->GetInputTensor("image");
  int h = inputs_batch_[0].new_im_size_[0];
  int w = inputs_batch_[0].new_im_size_[1];
  in_tensor->Reshape({batch_size, 3, h, w});
  std::vector<float> inputs_data(batch_size * 3 * h * w);
J
jack 已提交
226 227 228 229
  for (int i = 0; i < batch_size; ++i) {
    std::copy(inputs_batch_[i].im_data_.begin(),
              inputs_batch_[i].im_data_.end(),
              inputs_data.begin() + i * 3 * h * w);
J
jack 已提交
230 231
  }
  in_tensor->copy_from_cpu(inputs_data.data());
J
jack 已提交
232
  // in_tensor->copy_from_cpu(inputs_.im_data_.data());
J
jack 已提交
233
  predictor_->ZeroCopyRun();
S
syyxsxx 已提交
234
  // get result
J
jack 已提交
235 236 237 238 239 240 241 242 243
  auto output_names = predictor_->GetOutputNames();
  auto output_tensor = predictor_->GetOutputTensor(output_names[0]);
  std::vector<int> output_shape = output_tensor->shape();
  int size = 1;
  for (const auto& i : output_shape) {
    size *= i;
  }
  outputs_.resize(size);
  output_tensor->copy_to_cpu(outputs_.data());
S
syyxsxx 已提交
244
  // postprocess
245 246
  (*results).clear();
  (*results).resize(batch_size);
J
jack 已提交
247
  int single_batch_size = size / batch_size;
J
jack 已提交
248
  for (int i = 0; i < batch_size; ++i) {
J
jack 已提交
249 250 251 252 253
    auto start_ptr = std::begin(outputs_);
    auto end_ptr = std::begin(outputs_);
    std::advance(start_ptr, i * single_batch_size);
    std::advance(end_ptr, (i + 1) * single_batch_size);
    auto ptr = std::max_element(start_ptr, end_ptr);
J
jack 已提交
254 255 256
    (*results)[i].category_id = std::distance(start_ptr, ptr);
    (*results)[i].score = *ptr;
    (*results)[i].category = labels[(*results)[i].category_id];
J
jack 已提交
257 258 259 260
  }
  return true;
}

C
Channingss 已提交
261
bool Model::predict(const cv::Mat& im, DetResult* result) {
J
jack 已提交
262
  inputs_.clear();
C
Channingss 已提交
263 264 265
  result->clear();
  if (type == "classifier") {
    std::cerr << "Loading model is a 'classifier', ClsResult should be passed "
J
jack 已提交
266
                 "to function predict()!" << std::endl;
C
Channingss 已提交
267 268 269
    return false;
  } else if (type == "segmenter") {
    std::cerr << "Loading model is a 'segmenter', SegResult should be passed "
J
jack 已提交
270
                 "to function predict()!" << std::endl;
C
Channingss 已提交
271 272 273
    return false;
  }

S
syyxsxx 已提交
274
  // preprocess
C
Channingss 已提交
275 276 277 278 279 280 281 282 283 284
  if (!preprocess(im, &inputs_)) {
    std::cerr << "Preprocess failed!" << std::endl;
    return false;
  }

  int h = inputs_.new_im_size_[0];
  int w = inputs_.new_im_size_[1];
  auto im_tensor = predictor_->GetInputTensor("image");
  im_tensor->Reshape({1, 3, h, w});
  im_tensor->copy_from_cpu(inputs_.im_data_.data());
J
jack 已提交
285

C
Channingss 已提交
286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303
  if (name == "YOLOv3") {
    auto im_size_tensor = predictor_->GetInputTensor("im_size");
    im_size_tensor->Reshape({1, 2});
    im_size_tensor->copy_from_cpu(inputs_.ori_im_size_.data());
  } else if (name == "FasterRCNN" || name == "MaskRCNN") {
    auto im_info_tensor = predictor_->GetInputTensor("im_info");
    auto im_shape_tensor = predictor_->GetInputTensor("im_shape");
    im_info_tensor->Reshape({1, 3});
    im_shape_tensor->Reshape({1, 3});
    float ori_h = static_cast<float>(inputs_.ori_im_size_[0]);
    float ori_w = static_cast<float>(inputs_.ori_im_size_[1]);
    float new_h = static_cast<float>(inputs_.new_im_size_[0]);
    float new_w = static_cast<float>(inputs_.new_im_size_[1]);
    float im_info[] = {new_h, new_w, inputs_.scale};
    float im_shape[] = {ori_h, ori_w, 1.0};
    im_info_tensor->copy_from_cpu(im_info);
    im_shape_tensor->copy_from_cpu(im_shape);
  }
S
syyxsxx 已提交
304
  // predict
C
Channingss 已提交
305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321
  predictor_->ZeroCopyRun();

  std::vector<float> output_box;
  auto output_names = predictor_->GetOutputNames();
  auto output_box_tensor = predictor_->GetOutputTensor(output_names[0]);
  std::vector<int> output_box_shape = output_box_tensor->shape();
  int size = 1;
  for (const auto& i : output_box_shape) {
    size *= i;
  }
  output_box.resize(size);
  output_box_tensor->copy_to_cpu(output_box.data());
  if (size < 6) {
    std::cerr << "[WARNING] There's no object detected." << std::endl;
    return true;
  }
  int num_boxes = size / 6;
S
syyxsxx 已提交
322
  // box postprocess
C
Channingss 已提交
323 324 325 326 327 328 329 330 331 332 333 334 335 336
  for (int i = 0; i < num_boxes; ++i) {
    Box box;
    box.category_id = static_cast<int>(round(output_box[i * 6]));
    box.category = labels[box.category_id];
    box.score = output_box[i * 6 + 1];
    float xmin = output_box[i * 6 + 2];
    float ymin = output_box[i * 6 + 3];
    float xmax = output_box[i * 6 + 4];
    float ymax = output_box[i * 6 + 5];
    float w = xmax - xmin + 1;
    float h = ymax - ymin + 1;
    box.coordinate = {xmin, ymin, w, h};
    result->boxes.push_back(std::move(box));
  }
S
syyxsxx 已提交
337
  // mask postprocess
C
Channingss 已提交
338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354
  if (name == "MaskRCNN") {
    std::vector<float> output_mask;
    auto output_mask_tensor = predictor_->GetOutputTensor(output_names[1]);
    std::vector<int> output_mask_shape = output_mask_tensor->shape();
    int masks_size = 1;
    for (const auto& i : output_mask_shape) {
      masks_size *= i;
    }
    int mask_pixels = output_mask_shape[2] * output_mask_shape[3];
    int classes = output_mask_shape[1];
    output_mask.resize(masks_size);
    output_mask_tensor->copy_to_cpu(output_mask.data());
    result->mask_resolution = output_mask_shape[2];
    for (int i = 0; i < result->boxes.size(); ++i) {
      Box* box = &result->boxes[i];
      box->mask.shape = {static_cast<int>(box->coordinate[2]),
                         static_cast<int>(box->coordinate[3])};
S
syyxsxx 已提交
355
      auto begin_mask =
S
fix  
syyxsxx 已提交
356
          output_mask.data() + (i * classes + box->category_id) * mask_pixels;
S
syyxsxx 已提交
357 358 359 360 361 362 363 364
      cv::Mat bin_mask(result->mask_resolution,
                     result->mask_resolution,
                     CV_32FC1,
                     begin_mask);
      cv::resize(bin_mask,
               bin_mask,
               cv::Size(box->mask.shape[0], box->mask.shape[1]));
      cv::threshold(bin_mask, bin_mask, 0.5, 1, cv::THRESH_BINARY);
S
syyxsxx 已提交
365
      auto mask_int_begin = reinterpret_cast<float*>(bin_mask.data);
S
syyxsxx 已提交
366 367 368
      auto mask_int_end =
        mask_int_begin + box->mask.shape[0] * box->mask.shape[1];
      box->mask.data.assign(mask_int_begin, mask_int_end);
C
Channingss 已提交
369 370
    }
  }
J
jack 已提交
371
  return true;
C
Channingss 已提交
372 373
}

J
jack 已提交
374
bool Model::predict(const std::vector<cv::Mat>& im_batch,
375
                    std::vector<DetResult>* results,
J
jack 已提交
376 377
                    int thread_num) {
  for (auto& inputs : inputs_batch_) {
J
jack 已提交
378 379
    inputs.clear();
  }
J
jack 已提交
380 381
  if (type == "classifier") {
    std::cerr << "Loading model is a 'classifier', ClsResult should be passed "
J
jack 已提交
382
                 "to function predict()!" << std::endl;
J
jack 已提交
383 384 385
    return false;
  } else if (type == "segmenter") {
    std::cerr << "Loading model is a 'segmenter', SegResult should be passed "
J
jack 已提交
386
                 "to function predict()!" << std::endl;
J
jack 已提交
387 388 389
    return false;
  }

J
jack 已提交
390
  inputs_batch_.assign(im_batch.size(), ImageBlob());
J
jack 已提交
391
  int batch_size = im_batch.size();
S
syyxsxx 已提交
392
  // preprocess
J
jack 已提交
393
  if (!preprocess(im_batch, &inputs_batch_, thread_num)) {
J
jack 已提交
394 395 396
    std::cerr << "Preprocess failed!" << std::endl;
    return false;
  }
S
syyxsxx 已提交
397
  // RCNN model padding
J
jack 已提交
398 399 400 401
  if (batch_size > 1) {
    if (name == "FasterRCNN" || name == "MaskRCNN") {
      int max_h = -1;
      int max_w = -1;
J
jack 已提交
402
      for (int i = 0; i < batch_size; ++i) {
J
jack 已提交
403 404
        max_h = std::max(max_h, inputs_batch_[i].new_im_size_[0]);
        max_w = std::max(max_w, inputs_batch_[i].new_im_size_[1]);
J
jack 已提交
405 406
        // std::cout << "(" << inputs_batch_[i].new_im_size_[0]
        //          << ", " << inputs_batch_[i].new_im_size_[1]
J
jack 已提交
407
        //          <<  ")" << std::endl;
J
jack 已提交
408
      }
J
jack 已提交
409 410
      thread_num = std::min(thread_num, batch_size);
      #pragma omp parallel for num_threads(thread_num)
J
jack 已提交
411
      for (int i = 0; i < batch_size; ++i) {
J
jack 已提交
412 413 414
        int h = inputs_batch_[i].new_im_size_[0];
        int w = inputs_batch_[i].new_im_size_[1];
        int c = im_batch[i].channels();
J
jack 已提交
415
        if (max_h != h || max_w != w) {
J
jack 已提交
416
          std::vector<float> temp_buffer(c * max_h * max_w);
J
jack 已提交
417 418 419
          float* temp_ptr = temp_buffer.data();
          float* ptr = inputs_batch_[i].im_data_.data();
          for (int cur_channel = c - 1; cur_channel >= 0; --cur_channel) {
J
jack 已提交
420 421
            int ori_pos = cur_channel * h * w + (h - 1) * w;
            int des_pos = cur_channel * max_h * max_w + (h - 1) * max_w;
J
jack 已提交
422 423 424
            int last_pos = cur_channel * h * w;
            for (; ori_pos >= last_pos; ori_pos -= w, des_pos -= max_w) {
              memcpy(temp_ptr + des_pos, ptr + ori_pos, w * sizeof(float));
J
jack 已提交
425 426 427 428
            }
          }
          inputs_batch_[i].im_data_.swap(temp_buffer);
          inputs_batch_[i].new_im_size_[0] = max_h;
J
jack 已提交
429
          inputs_batch_[i].new_im_size_[1] = max_w;
J
jack 已提交
430 431 432 433
        }
      }
    }
  }
J
jack 已提交
434 435 436 437 438
  int h = inputs_batch_[0].new_im_size_[0];
  int w = inputs_batch_[0].new_im_size_[1];
  auto im_tensor = predictor_->GetInputTensor("image");
  im_tensor->Reshape({batch_size, 3, h, w});
  std::vector<float> inputs_data(batch_size * 3 * h * w);
J
jack 已提交
439 440 441 442
  for (int i = 0; i < batch_size; ++i) {
    std::copy(inputs_batch_[i].im_data_.begin(),
              inputs_batch_[i].im_data_.end(),
              inputs_data.begin() + i * 3 * h * w);
J
jack 已提交
443 444 445 446 447
  }
  im_tensor->copy_from_cpu(inputs_data.data());
  if (name == "YOLOv3") {
    auto im_size_tensor = predictor_->GetInputTensor("im_size");
    im_size_tensor->Reshape({batch_size, 2});
J
jack 已提交
448 449 450 451 452
    std::vector<int> inputs_data_size(batch_size * 2);
    for (int i = 0; i < batch_size; ++i) {
      std::copy(inputs_batch_[i].ori_im_size_.begin(),
                inputs_batch_[i].ori_im_size_.end(),
                inputs_data_size.begin() + 2 * i);
J
jack 已提交
453 454 455 456 457 458 459
    }
    im_size_tensor->copy_from_cpu(inputs_data_size.data());
  } else if (name == "FasterRCNN" || name == "MaskRCNN") {
    auto im_info_tensor = predictor_->GetInputTensor("im_info");
    auto im_shape_tensor = predictor_->GetInputTensor("im_shape");
    im_info_tensor->Reshape({batch_size, 3});
    im_shape_tensor->Reshape({batch_size, 3});
J
jack 已提交
460

J
jack 已提交
461 462
    std::vector<float> im_info(3 * batch_size);
    std::vector<float> im_shape(3 * batch_size);
J
jack 已提交
463
    for (int i = 0; i < batch_size; ++i) {
J
jack 已提交
464 465 466 467 468 469 470 471 472 473 474 475 476 477
      float ori_h = static_cast<float>(inputs_batch_[i].ori_im_size_[0]);
      float ori_w = static_cast<float>(inputs_batch_[i].ori_im_size_[1]);
      float new_h = static_cast<float>(inputs_batch_[i].new_im_size_[0]);
      float new_w = static_cast<float>(inputs_batch_[i].new_im_size_[1]);
      im_info[i * 3] = new_h;
      im_info[i * 3 + 1] = new_w;
      im_info[i * 3 + 2] = inputs_batch_[i].scale;
      im_shape[i * 3] = ori_h;
      im_shape[i * 3 + 1] = ori_w;
      im_shape[i * 3 + 2] = 1.0;
    }
    im_info_tensor->copy_from_cpu(im_info.data());
    im_shape_tensor->copy_from_cpu(im_shape.data());
  }
S
syyxsxx 已提交
478
  // predict
J
jack 已提交
479 480
  predictor_->ZeroCopyRun();

S
syyxsxx 已提交
481
  // get all box
J
jack 已提交
482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497
  std::vector<float> output_box;
  auto output_names = predictor_->GetOutputNames();
  auto output_box_tensor = predictor_->GetOutputTensor(output_names[0]);
  std::vector<int> output_box_shape = output_box_tensor->shape();
  int size = 1;
  for (const auto& i : output_box_shape) {
    size *= i;
  }
  output_box.resize(size);
  output_box_tensor->copy_to_cpu(output_box.data());
  if (size < 6) {
    std::cerr << "[WARNING] There's no object detected." << std::endl;
    return true;
  }
  auto lod_vector = output_box_tensor->lod();
  int num_boxes = size / 6;
S
syyxsxx 已提交
498
  // box postprocess
499 500
  (*results).clear();
  (*results).resize(batch_size);
J
jack 已提交
501
  for (int i = 0; i < lod_vector[0].size() - 1; ++i) {
J
jack 已提交
502
    for (int j = lod_vector[0][i]; j < lod_vector[0][i + 1]; ++j) {
J
jack 已提交
503
      Box box;
J
jack 已提交
504
      box.category_id = static_cast<int>(round(output_box[j * 6]));
J
jack 已提交
505 506 507 508 509 510 511 512 513
      box.category = labels[box.category_id];
      box.score = output_box[j * 6 + 1];
      float xmin = output_box[j * 6 + 2];
      float ymin = output_box[j * 6 + 3];
      float xmax = output_box[j * 6 + 4];
      float ymax = output_box[j * 6 + 5];
      float w = xmax - xmin + 1;
      float h = ymax - ymin + 1;
      box.coordinate = {xmin, ymin, w, h};
514
      (*results)[i].boxes.push_back(std::move(box));
J
jack 已提交
515 516 517
    }
  }

S
syyxsxx 已提交
518
  // mask postprocess
J
jack 已提交
519 520 521 522 523 524 525 526 527 528 529 530 531
  if (name == "MaskRCNN") {
    std::vector<float> output_mask;
    auto output_mask_tensor = predictor_->GetOutputTensor(output_names[1]);
    std::vector<int> output_mask_shape = output_mask_tensor->shape();
    int masks_size = 1;
    for (const auto& i : output_mask_shape) {
      masks_size *= i;
    }
    int mask_pixels = output_mask_shape[2] * output_mask_shape[3];
    int classes = output_mask_shape[1];
    output_mask.resize(masks_size);
    output_mask_tensor->copy_to_cpu(output_mask.data());
    int mask_idx = 0;
J
jack 已提交
532
    for (int i = 0; i < lod_vector[0].size() - 1; ++i) {
533 534
      (*results)[i].mask_resolution = output_mask_shape[2];
      for (int j = 0; j < (*results)[i].boxes.size(); ++j) {
S
fix  
syyxsxx 已提交
535
        Box* box = &(*results)[i].boxes[i];
J
jack 已提交
536
        int category_id = box->category_id;
S
syyxsxx 已提交
537 538
        box->mask.shape = {static_cast<int>(box->coordinate[2]),
                          static_cast<int>(box->coordinate[3])};
S
syyxsxx 已提交
539
        auto begin_mask =
S
fix  
syyxsxx 已提交
540
          output_mask.data() + (i * classes + box->category_id) * mask_pixels;
S
fix  
syyxsxx 已提交
541 542
        cv::Mat bin_mask(output_mask_shape[2],
                      output_mask_shape[2],
S
syyxsxx 已提交
543 544 545 546 547 548
                      CV_32FC1,
                      begin_mask);
        cv::resize(bin_mask,
                bin_mask,
                cv::Size(box->mask.shape[0], box->mask.shape[1]));
        cv::threshold(bin_mask, bin_mask, 0.5, 1, cv::THRESH_BINARY);
S
syyxsxx 已提交
549
        auto mask_int_begin = reinterpret_cast<float*>(bin_mask.data);
S
syyxsxx 已提交
550 551 552
        auto mask_int_end =
          mask_int_begin + box->mask.shape[0] * box->mask.shape[1];
        box->mask.data.assign(mask_int_begin, mask_int_end);
J
jack 已提交
553 554 555 556
        mask_idx++;
      }
    }
  }
J
jack 已提交
557
  return true;
J
jack 已提交
558 559
}

C
Channingss 已提交
560 561 562 563 564
bool Model::predict(const cv::Mat& im, SegResult* result) {
  result->clear();
  inputs_.clear();
  if (type == "classifier") {
    std::cerr << "Loading model is a 'classifier', ClsResult should be passed "
J
jack 已提交
565
                 "to function predict()!" << std::endl;
C
Channingss 已提交
566 567 568
    return false;
  } else if (type == "detector") {
    std::cerr << "Loading model is a 'detector', DetResult should be passed to "
J
jack 已提交
569
                 "function predict()!" << std::endl;
C
Channingss 已提交
570 571 572
    return false;
  }

S
syyxsxx 已提交
573
  // preprocess
C
Channingss 已提交
574 575 576 577 578 579 580 581 582 583 584
  if (!preprocess(im, &inputs_)) {
    std::cerr << "Preprocess failed!" << std::endl;
    return false;
  }

  int h = inputs_.new_im_size_[0];
  int w = inputs_.new_im_size_[1];
  auto im_tensor = predictor_->GetInputTensor("image");
  im_tensor->Reshape({1, 3, h, w});
  im_tensor->copy_from_cpu(inputs_.im_data_.data());

S
syyxsxx 已提交
585
  // predict
C
Channingss 已提交
586 587
  predictor_->ZeroCopyRun();

S
syyxsxx 已提交
588
  // get labelmap
C
Channingss 已提交
589 590 591 592 593 594 595 596
  auto output_names = predictor_->GetOutputNames();
  auto output_label_tensor = predictor_->GetOutputTensor(output_names[0]);
  std::vector<int> output_label_shape = output_label_tensor->shape();
  int size = 1;
  for (const auto& i : output_label_shape) {
    size *= i;
    result->label_map.shape.push_back(i);
  }
J
jack 已提交
597

C
Channingss 已提交
598 599 600
  result->label_map.data.resize(size);
  output_label_tensor->copy_to_cpu(result->label_map.data.data());

S
syyxsxx 已提交
601
  // get scoremap
C
Channingss 已提交
602 603 604 605 606 607 608
  auto output_score_tensor = predictor_->GetOutputTensor(output_names[1]);
  std::vector<int> output_score_shape = output_score_tensor->shape();
  size = 1;
  for (const auto& i : output_score_shape) {
    size *= i;
    result->score_map.shape.push_back(i);
  }
J
jack 已提交
609

C
Channingss 已提交
610 611 612
  result->score_map.data.resize(size);
  output_score_tensor->copy_to_cpu(result->score_map.data.data());

S
syyxsxx 已提交
613
  // get origin image result
C
Channingss 已提交
614 615 616 617 618 619 620 621 622 623 624
  std::vector<uint8_t> label_map(result->label_map.data.begin(),
                                 result->label_map.data.end());
  cv::Mat mask_label(result->label_map.shape[1],
                     result->label_map.shape[2],
                     CV_8UC1,
                     label_map.data());

  cv::Mat mask_score(result->score_map.shape[2],
                     result->score_map.shape[3],
                     CV_32FC1,
                     result->score_map.data.data());
C
Channingss 已提交
625
  int idx = 1;
C
Channingss 已提交
626
  int len_postprocess = inputs_.im_size_before_resize_.size();
C
Channingss 已提交
627 628
  for (std::vector<std::string>::reverse_iterator iter =
           inputs_.reshape_order_.rbegin();
C
Channingss 已提交
629 630
       iter != inputs_.reshape_order_.rend();
       ++iter) {
C
Channingss 已提交
631
    if (*iter == "padding") {
C
Channingss 已提交
632
      auto before_shape = inputs_.im_size_before_resize_[len_postprocess - idx];
C
Channingss 已提交
633 634 635
      inputs_.im_size_before_resize_.pop_back();
      auto padding_w = before_shape[0];
      auto padding_h = before_shape[1];
J
jack 已提交
636 637
      mask_label = mask_label(cv::Rect(0, 0, padding_h, padding_w));
      mask_score = mask_score(cv::Rect(0, 0, padding_h, padding_w));
C
Channingss 已提交
638
    } else if (*iter == "resize") {
C
Channingss 已提交
639
      auto before_shape = inputs_.im_size_before_resize_[len_postprocess - idx];
C
Channingss 已提交
640 641 642
      inputs_.im_size_before_resize_.pop_back();
      auto resize_w = before_shape[0];
      auto resize_h = before_shape[1];
C
Channingss 已提交
643 644 645 646 647 648 649 650 651 652 653
      cv::resize(mask_label,
                 mask_label,
                 cv::Size(resize_h, resize_w),
                 0,
                 0,
                 cv::INTER_NEAREST);
      cv::resize(mask_score,
                 mask_score,
                 cv::Size(resize_h, resize_w),
                 0,
                 0,
J
jack 已提交
654
                 cv::INTER_LINEAR);
C
Channingss 已提交
655
    }
C
Channingss 已提交
656
    ++idx;
C
Channingss 已提交
657 658 659 660 661 662 663
  }
  result->label_map.data.assign(mask_label.begin<uint8_t>(),
                                mask_label.end<uint8_t>());
  result->label_map.shape = {mask_label.rows, mask_label.cols};
  result->score_map.data.assign(mask_score.begin<float>(),
                                mask_score.end<float>());
  result->score_map.shape = {mask_score.rows, mask_score.cols};
J
jack 已提交
664 665 666
  return true;
}

J
jack 已提交
667
bool Model::predict(const std::vector<cv::Mat>& im_batch,
668
                    std::vector<SegResult>* results,
J
jack 已提交
669 670
                    int thread_num) {
  for (auto& inputs : inputs_batch_) {
J
jack 已提交
671 672 673 674
    inputs.clear();
  }
  if (type == "classifier") {
    std::cerr << "Loading model is a 'classifier', ClsResult should be passed "
J
jack 已提交
675
                 "to function predict()!" << std::endl;
J
jack 已提交
676 677 678
    return false;
  } else if (type == "detector") {
    std::cerr << "Loading model is a 'detector', DetResult should be passed to "
J
jack 已提交
679
                 "function predict()!" << std::endl;
J
jack 已提交
680 681 682
    return false;
  }

S
syyxsxx 已提交
683
  // preprocess
J
jack 已提交
684
  inputs_batch_.assign(im_batch.size(), ImageBlob());
J
jack 已提交
685
  if (!preprocess(im_batch, &inputs_batch_, thread_num)) {
J
jack 已提交
686 687 688 689 690
    std::cerr << "Preprocess failed!" << std::endl;
    return false;
  }

  int batch_size = im_batch.size();
691 692
  (*results).clear();
  (*results).resize(batch_size);
J
jack 已提交
693 694 695 696 697
  int h = inputs_batch_[0].new_im_size_[0];
  int w = inputs_batch_[0].new_im_size_[1];
  auto im_tensor = predictor_->GetInputTensor("image");
  im_tensor->Reshape({batch_size, 3, h, w});
  std::vector<float> inputs_data(batch_size * 3 * h * w);
J
jack 已提交
698 699 700 701
  for (int i = 0; i < batch_size; ++i) {
    std::copy(inputs_batch_[i].im_data_.begin(),
              inputs_batch_[i].im_data_.end(),
              inputs_data.begin() + i * 3 * h * w);
J
jack 已提交
702 703
  }
  im_tensor->copy_from_cpu(inputs_data.data());
J
jack 已提交
704
  // im_tensor->copy_from_cpu(inputs_.im_data_.data());
J
jack 已提交
705

S
syyxsxx 已提交
706
  // predict
J
jack 已提交
707 708
  predictor_->ZeroCopyRun();

S
syyxsxx 已提交
709
  // get labelmap
J
jack 已提交
710 711 712 713 714 715 716 717 718 719 720 721 722
  auto output_names = predictor_->GetOutputNames();
  auto output_label_tensor = predictor_->GetOutputTensor(output_names[0]);
  std::vector<int> output_label_shape = output_label_tensor->shape();
  int size = 1;
  for (const auto& i : output_label_shape) {
    size *= i;
  }

  std::vector<int64_t> output_labels(size, 0);
  output_label_tensor->copy_to_cpu(output_labels.data());
  auto output_labels_iter = output_labels.begin();

  int single_batch_size = size / batch_size;
J
jack 已提交
723
  for (int i = 0; i < batch_size; ++i) {
724 725
    (*results)[i].label_map.data.resize(single_batch_size);
    (*results)[i].label_map.shape.push_back(1);
J
jack 已提交
726
    for (int j = 1; j < output_label_shape.size(); ++j) {
727
      (*results)[i].label_map.shape.push_back(output_label_shape[j]);
J
jack 已提交
728
    }
J
jack 已提交
729 730
    std::copy(output_labels_iter + i * single_batch_size,
              output_labels_iter + (i + 1) * single_batch_size,
731
              (*results)[i].label_map.data.data());
J
jack 已提交
732 733
  }

S
syyxsxx 已提交
734
  // get scoremap
J
jack 已提交
735 736 737 738 739 740 741 742 743 744 745 746
  auto output_score_tensor = predictor_->GetOutputTensor(output_names[1]);
  std::vector<int> output_score_shape = output_score_tensor->shape();
  size = 1;
  for (const auto& i : output_score_shape) {
    size *= i;
  }

  std::vector<float> output_scores(size, 0);
  output_score_tensor->copy_to_cpu(output_scores.data());
  auto output_scores_iter = output_scores.begin();

  int single_batch_score_size = size / batch_size;
J
jack 已提交
747
  for (int i = 0; i < batch_size; ++i) {
748 749
    (*results)[i].score_map.data.resize(single_batch_score_size);
    (*results)[i].score_map.shape.push_back(1);
J
jack 已提交
750
    for (int j = 1; j < output_score_shape.size(); ++j) {
751
      (*results)[i].score_map.shape.push_back(output_score_shape[j]);
J
jack 已提交
752
    }
J
jack 已提交
753 754
    std::copy(output_scores_iter + i * single_batch_score_size,
              output_scores_iter + (i + 1) * single_batch_score_size,
755
              (*results)[i].score_map.data.data());
J
jack 已提交
756 757
  }

S
syyxsxx 已提交
758
  // get origin image result
J
jack 已提交
759
  for (int i = 0; i < batch_size; ++i) {
760 761 762 763
    std::vector<uint8_t> label_map((*results)[i].label_map.data.begin(),
                                   (*results)[i].label_map.data.end());
    cv::Mat mask_label((*results)[i].label_map.shape[1],
                       (*results)[i].label_map.shape[2],
J
jack 已提交
764 765
                       CV_8UC1,
                       label_map.data());
J
jack 已提交
766

767 768
    cv::Mat mask_score((*results)[i].score_map.shape[2],
                       (*results)[i].score_map.shape[3],
J
jack 已提交
769
                       CV_32FC1,
770
                       (*results)[i].score_map.data.data());
J
jack 已提交
771 772 773 774 775 776 777
    int idx = 1;
    int len_postprocess = inputs_batch_[i].im_size_before_resize_.size();
    for (std::vector<std::string>::reverse_iterator iter =
             inputs_batch_[i].reshape_order_.rbegin();
         iter != inputs_batch_[i].reshape_order_.rend();
         ++iter) {
      if (*iter == "padding") {
J
jack 已提交
778 779
        auto before_shape =
            inputs_batch_[i].im_size_before_resize_[len_postprocess - idx];
J
jack 已提交
780 781 782 783 784 785
        inputs_batch_[i].im_size_before_resize_.pop_back();
        auto padding_w = before_shape[0];
        auto padding_h = before_shape[1];
        mask_label = mask_label(cv::Rect(0, 0, padding_h, padding_w));
        mask_score = mask_score(cv::Rect(0, 0, padding_h, padding_w));
      } else if (*iter == "resize") {
J
jack 已提交
786 787
        auto before_shape =
            inputs_batch_[i].im_size_before_resize_[len_postprocess - idx];
J
jack 已提交
788 789 790 791 792 793 794 795 796 797 798 799 800 801
        inputs_batch_[i].im_size_before_resize_.pop_back();
        auto resize_w = before_shape[0];
        auto resize_h = before_shape[1];
        cv::resize(mask_label,
                   mask_label,
                   cv::Size(resize_h, resize_w),
                   0,
                   0,
                   cv::INTER_NEAREST);
        cv::resize(mask_score,
                   mask_score,
                   cv::Size(resize_h, resize_w),
                   0,
                   0,
J
jack 已提交
802
                   cv::INTER_LINEAR);
J
jack 已提交
803 804 805
      }
      ++idx;
    }
806
    (*results)[i].label_map.data.assign(mask_label.begin<uint8_t>(),
J
jack 已提交
807
                                       mask_label.end<uint8_t>());
808 809
    (*results)[i].label_map.shape = {mask_label.rows, mask_label.cols};
    (*results)[i].score_map.data.assign(mask_score.begin<float>(),
J
jack 已提交
810
                                       mask_score.end<float>());
811
    (*results)[i].score_map.shape = {mask_score.rows, mask_score.cols};
J
jack 已提交
812 813
  }
  return true;
C
Channingss 已提交
814 815
}

J
jack 已提交
816
}  // namespace PaddleX