ocr_rec.cpp 6.2 KB
Newer Older
littletomatodonkey's avatar
littletomatodonkey 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
// 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.

#include <include/ocr_rec.h>

namespace PaddleOCR {

void CRNNRecognizer::Run(std::vector<std::vector<std::vector<int>>> boxes,
W
WenmuZhou 已提交
20
                         cv::Mat &img, Classifier *cls) {
littletomatodonkey's avatar
littletomatodonkey 已提交
21 22 23 24 25 26 27 28
  cv::Mat srcimg;
  img.copyTo(srcimg);
  cv::Mat crop_img;
  cv::Mat resize_img;

  std::cout << "The predicted text is :" << std::endl;
  int index = 0;
  for (int i = boxes.size() - 1; i >= 0; i--) {
littletomatodonkey's avatar
littletomatodonkey 已提交
29
    crop_img = GetRotateCropImage(srcimg, boxes[i]);
W
WenmuZhou 已提交
30 31 32
    if (cls != nullptr) {
      crop_img = cls->Run(crop_img);
    }
littletomatodonkey's avatar
littletomatodonkey 已提交
33 34 35 36 37 38 39 40

    float wh_ratio = float(crop_img.cols) / float(crop_img.rows);

    this->resize_op_.Run(crop_img, resize_img, wh_ratio);

    this->normalize_op_.Run(&resize_img, this->mean_, this->scale_,
                            this->is_scale_);

littletomatodonkey's avatar
littletomatodonkey 已提交
41
    std::vector<float> input(1 * 3 * resize_img.rows * resize_img.cols, 0.0f);
littletomatodonkey's avatar
littletomatodonkey 已提交
42

littletomatodonkey's avatar
littletomatodonkey 已提交
43
    this->permute_op_.Run(&resize_img, input.data());
littletomatodonkey's avatar
littletomatodonkey 已提交
44

45
    // Inference.
L
LDOUBLEV 已提交
46 47 48 49 50
    auto input_names = this->predictor_->GetInputNames();
    auto input_t = this->predictor_->GetInputHandle(input_names[0]);
    input_t->Reshape({1, 3, resize_img.rows, resize_img.cols});
    input_t->CopyFromCpu(input.data());
    this->predictor_->Run();
littletomatodonkey's avatar
littletomatodonkey 已提交
51

W
WenmuZhou 已提交
52
    std::vector<float> predict_batch;
littletomatodonkey's avatar
littletomatodonkey 已提交
53
    auto output_names = this->predictor_->GetOutputNames();
L
LDOUBLEV 已提交
54
    auto output_t = this->predictor_->GetOutputHandle(output_names[0]);
W
WenmuZhou 已提交
55
    auto predict_shape = output_t->shape();
56

W
WenmuZhou 已提交
57
    int out_num = std::accumulate(predict_shape.begin(), predict_shape.end(), 1,
littletomatodonkey's avatar
littletomatodonkey 已提交
58
                                  std::multiplies<int>());
W
WenmuZhou 已提交
59
    predict_batch.resize(out_num);
littletomatodonkey's avatar
littletomatodonkey 已提交
60

L
LDOUBLEV 已提交
61
    output_t->CopyToCpu(predict_batch.data());
littletomatodonkey's avatar
littletomatodonkey 已提交
62

W
WenmuZhou 已提交
63 64
    // ctc decode
    std::vector<std::string> str_res;
littletomatodonkey's avatar
littletomatodonkey 已提交
65
    int argmax_idx;
W
WenmuZhou 已提交
66
    int last_index = 0;
littletomatodonkey's avatar
littletomatodonkey 已提交
67 68 69 70
    float score = 0.f;
    int count = 0;
    float max_value = 0.0f;

W
WenmuZhou 已提交
71
    for (int n = 0; n < predict_shape[1]; n++) {
littletomatodonkey's avatar
littletomatodonkey 已提交
72
      argmax_idx =
W
WenmuZhou 已提交
73 74
          int(Utility::argmax(&predict_batch[n * predict_shape[2]],
                              &predict_batch[(n + 1) * predict_shape[2]]));
littletomatodonkey's avatar
littletomatodonkey 已提交
75
      max_value =
W
WenmuZhou 已提交
76 77 78 79
          float(*std::max_element(&predict_batch[n * predict_shape[2]],
                                  &predict_batch[(n + 1) * predict_shape[2]]));

      if (argmax_idx > 0 && (not(i > 0 && argmax_idx == last_index))) {
littletomatodonkey's avatar
littletomatodonkey 已提交
80 81
        score += max_value;
        count += 1;
W
WenmuZhou 已提交
82
        str_res.push_back(label_list_[argmax_idx]);
littletomatodonkey's avatar
littletomatodonkey 已提交
83
      }
W
WenmuZhou 已提交
84
      last_index = argmax_idx;
littletomatodonkey's avatar
littletomatodonkey 已提交
85 86
    }
    score /= count;
W
WenmuZhou 已提交
87 88 89
    for (int i = 0; i < str_res.size(); i++) {
      std::cout << str_res[i];
    }
littletomatodonkey's avatar
littletomatodonkey 已提交
90 91 92 93
    std::cout << "\tscore: " << score << std::endl;
  }
}

littletomatodonkey's avatar
littletomatodonkey 已提交
94
void CRNNRecognizer::LoadModel(const std::string &model_dir) {
L
LDOUBLEV 已提交
95 96
  //   AnalysisConfig config;
  paddle_infer::Config config;
文幕地方's avatar
文幕地方 已提交
97 98
  config.SetModel(model_dir + "/inference.pdmodel",
                  model_dir + "/inference.pdiparams");
littletomatodonkey's avatar
littletomatodonkey 已提交
99

littletomatodonkey's avatar
littletomatodonkey 已提交
100 101
  if (this->use_gpu_) {
    config.EnableUseGpu(this->gpu_mem_, this->gpu_id_);
L
LDOUBLEV 已提交
102 103 104 105 106 107 108
    if (this->use_tensorrt_) {
      config.EnableTensorRtEngine(
          1 << 20, 10, 3,
          this->use_fp16_ ? paddle_infer::Config::Precision::kHalf
                          : paddle_infer::Config::Precision::kFloat32,
          false, false);
    }
littletomatodonkey's avatar
littletomatodonkey 已提交
109 110
  } else {
    config.DisableGpu();
littletomatodonkey's avatar
littletomatodonkey 已提交
111 112
    if (this->use_mkldnn_) {
      config.EnableMKLDNN();
W
WenmuZhou 已提交
113 114
      // cache 10 different shapes for mkldnn to avoid memory leak
      config.SetMkldnnCacheCapacity(10);
littletomatodonkey's avatar
littletomatodonkey 已提交
115
    }
littletomatodonkey's avatar
littletomatodonkey 已提交
116 117
    config.SetCpuMathLibraryNumThreads(this->cpu_math_library_num_threads_);
  }
littletomatodonkey's avatar
littletomatodonkey 已提交
118

L
LDOUBLEV 已提交
119
  config.SwitchUseFeedFetchOps(false);
littletomatodonkey's avatar
littletomatodonkey 已提交
120
  // true for multiple input
littletomatodonkey's avatar
littletomatodonkey 已提交
121
  config.SwitchSpecifyInputNames(true);
littletomatodonkey's avatar
littletomatodonkey 已提交
122 123 124 125

  config.SwitchIrOptim(true);

  config.EnableMemoryOptim();
littletomatodonkey's avatar
littletomatodonkey 已提交
126
  config.DisableGlogInfo();
littletomatodonkey's avatar
littletomatodonkey 已提交
127

L
LDOUBLEV 已提交
128
  this->predictor_ = CreatePredictor(config);
littletomatodonkey's avatar
littletomatodonkey 已提交
129 130
}

littletomatodonkey's avatar
littletomatodonkey 已提交
131 132
cv::Mat CRNNRecognizer::GetRotateCropImage(const cv::Mat &srcimage,
                                           std::vector<std::vector<int>> box) {
littletomatodonkey's avatar
littletomatodonkey 已提交
133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185
  cv::Mat image;
  srcimage.copyTo(image);
  std::vector<std::vector<int>> points = box;

  int x_collect[4] = {box[0][0], box[1][0], box[2][0], box[3][0]};
  int y_collect[4] = {box[0][1], box[1][1], box[2][1], box[3][1]};
  int left = int(*std::min_element(x_collect, x_collect + 4));
  int right = int(*std::max_element(x_collect, x_collect + 4));
  int top = int(*std::min_element(y_collect, y_collect + 4));
  int bottom = int(*std::max_element(y_collect, y_collect + 4));

  cv::Mat img_crop;
  image(cv::Rect(left, top, right - left, bottom - top)).copyTo(img_crop);

  for (int i = 0; i < points.size(); i++) {
    points[i][0] -= left;
    points[i][1] -= top;
  }

  int img_crop_width = int(sqrt(pow(points[0][0] - points[1][0], 2) +
                                pow(points[0][1] - points[1][1], 2)));
  int img_crop_height = int(sqrt(pow(points[0][0] - points[3][0], 2) +
                                 pow(points[0][1] - points[3][1], 2)));

  cv::Point2f pts_std[4];
  pts_std[0] = cv::Point2f(0., 0.);
  pts_std[1] = cv::Point2f(img_crop_width, 0.);
  pts_std[2] = cv::Point2f(img_crop_width, img_crop_height);
  pts_std[3] = cv::Point2f(0.f, img_crop_height);

  cv::Point2f pointsf[4];
  pointsf[0] = cv::Point2f(points[0][0], points[0][1]);
  pointsf[1] = cv::Point2f(points[1][0], points[1][1]);
  pointsf[2] = cv::Point2f(points[2][0], points[2][1]);
  pointsf[3] = cv::Point2f(points[3][0], points[3][1]);

  cv::Mat M = cv::getPerspectiveTransform(pointsf, pts_std);

  cv::Mat dst_img;
  cv::warpPerspective(img_crop, dst_img, M,
                      cv::Size(img_crop_width, img_crop_height),
                      cv::BORDER_REPLICATE);

  if (float(dst_img.rows) >= float(dst_img.cols) * 1.5) {
    cv::Mat srcCopy = cv::Mat(dst_img.rows, dst_img.cols, dst_img.depth());
    cv::transpose(dst_img, srcCopy);
    cv::flip(srcCopy, srcCopy, 0);
    return srcCopy;
  } else {
    return dst_img;
  }
}

W
WenmuZhou 已提交
186
} // namespace PaddleOCR