first commit

third_party/paddle_ocr/src/paddleocr.cpp

@@ -0,0 +1,169 @@
+// 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/paddleocr.h"
+#include "../auto_log/autolog.h"
+
+#include <numeric>
+namespace PaddleOCR {
+
+PPOCR::PPOCR(std::string det_model_dir, 
+                std::string cls_model_dir, 
+                std::string rec_model_dir,
+                std::string rec_char_dict_path) {
+  if (!det_model_dir.empty()) {
+    this->detector_ = new DBDetector(det_model_dir);
+  }
+  if (!cls_model_dir.empty()) {
+    this->classifier_ = new Classifier(cls_model_dir);
+  }
+  if (!rec_model_dir.empty() && !rec_char_dict_path.empty()) {
+    this->recognizer_ = new CRNNRecognizer(rec_model_dir, rec_char_dict_path);
+  }
+};
+
+void PPOCR::det(cv::Mat img, std::vector<OCRPredictResult> &ocr_results,
+                std::vector<double> &times) {
+  std::vector<std::vector<std::vector<int>>> boxes;
+  std::vector<double> det_times;
+
+  this->detector_->Run(img, boxes, det_times);
+
+  for (int i = 0; i < boxes.size(); i++) {
+    OCRPredictResult res;
+    res.box = boxes[i];
+    ocr_results.push_back(res);
+  }
+  // sort boex from top to bottom, from left to right
+  Utility::sorted_boxes(ocr_results);
+  times[0] += det_times[0];
+  times[1] += det_times[1];
+  times[2] += det_times[2];
+}
+
+void PPOCR::rec(std::vector<cv::Mat> img_list,
+                std::vector<OCRPredictResult> &ocr_results,
+                std::vector<double> &times) {
+  std::vector<std::string> rec_texts(img_list.size(), "");
+  std::vector<float> rec_text_scores(img_list.size(), 0);
+  std::vector<double> rec_times;
+  this->recognizer_->Run(img_list, rec_texts, rec_text_scores, rec_times);
+  // output rec results
+  for (int i = 0; i < rec_texts.size(); i++) {
+    ocr_results[i].text = rec_texts[i];
+    ocr_results[i].score = rec_text_scores[i];
+  }
+  times[0] += rec_times[0];
+  times[1] += rec_times[1];
+  times[2] += rec_times[2];
+}
+
+void PPOCR::cls(std::vector<cv::Mat> img_list,
+                std::vector<OCRPredictResult> &ocr_results,
+                std::vector<double> &times) {
+  std::vector<int> cls_labels(img_list.size(), 0);
+  std::vector<float> cls_scores(img_list.size(), 0);
+  std::vector<double> cls_times;
+  this->classifier_->Run(img_list, cls_labels, cls_scores, cls_times);
+  // output cls results
+  for (int i = 0; i < cls_labels.size(); i++) {
+    ocr_results[i].cls_label = cls_labels[i];
+    ocr_results[i].cls_score = cls_scores[i];
+  }
+  times[0] += cls_times[0];
+  times[1] += cls_times[1];
+  times[2] += cls_times[2];
+}
+
+std::vector<std::vector<OCRPredictResult>>
+PPOCR::ocr(std::vector<cv::Mat>& cv_all_imgs, bool det, bool rec,
+           bool cls) {
+  std::vector<double> time_info_det = {0, 0, 0};
+  std::vector<double> time_info_rec = {0, 0, 0};
+  std::vector<double> time_info_cls = {0, 0, 0};
+  std::vector<std::vector<OCRPredictResult>> ocr_results;
+
+  if (!det) {
+    std::vector<OCRPredictResult> ocr_result;
+    // read image
+    auto& img_list = cv_all_imgs;
+    for (int i = 0; i < cv_all_imgs.size(); ++i) {
+      OCRPredictResult res;
+      ocr_result.push_back(res);
+    }
+    if (cls && this->classifier_ != nullptr) {
+      this->cls(img_list, ocr_result, time_info_cls);
+      for (int i = 0; i < img_list.size(); i++) {
+        if (ocr_result[i].cls_label % 2 == 1 &&
+            ocr_result[i].cls_score > this->classifier_->cls_thresh) {
+          cv::rotate(img_list[i], img_list[i], 1);
+        }
+      }
+    }
+    if (rec) {
+      this->rec(img_list, ocr_result, time_info_rec);
+    }
+    for (int i = 0; i < cv_all_imgs.size(); ++i) {
+      std::vector<OCRPredictResult> ocr_result_tmp;
+      ocr_result_tmp.push_back(ocr_result[i]);
+      ocr_results.push_back(ocr_result_tmp);
+    }
+  } else {
+    for (int i = 0; i < cv_all_imgs.size(); ++i) {
+      std::vector<OCRPredictResult> ocr_result;
+      cv::Mat& srcimg = cv_all_imgs[i];
+
+      // det
+      this->det(srcimg, ocr_result, time_info_det);
+      // crop image
+      std::vector<cv::Mat> img_list;
+      for (int j = 0; j < ocr_result.size(); j++) {
+        cv::Mat crop_img;
+        crop_img = Utility::GetRotateCropImage(srcimg, ocr_result[j].box);
+        img_list.push_back(crop_img);
+      }
+
+      // cls
+      if (cls && this->classifier_ != nullptr) {
+        this->cls(img_list, ocr_result, time_info_cls);
+        for (int i = 0; i < img_list.size(); i++) {
+          if (ocr_result[i].cls_label % 2 == 1 &&
+              ocr_result[i].cls_score > this->classifier_->cls_thresh) {
+            cv::rotate(img_list[i], img_list[i], 1);
+          }
+        }
+      }
+      // rec
+      if (rec) {
+        this->rec(img_list, ocr_result, time_info_rec);
+      }
+      ocr_results.push_back(ocr_result);
+    }
+  }
+  return ocr_results;
+} // namespace PaddleOCR
+
+PPOCR::~PPOCR() {
+  if (this->detector_ != nullptr) {
+    delete this->detector_;
+  }
+  if (this->classifier_ != nullptr) {
+    delete this->classifier_;
+  }
+  if (this->recognizer_ != nullptr) {
+    delete this->recognizer_;
+  }
+};
+
+} // namespace PaddleOCR