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

#pragma once

#include "include/clipper.h"
littletomatodonkey's avatar
littletomatodonkey 已提交
18
#include "include/utility.h"
littletomatodonkey's avatar
littletomatodonkey 已提交
19 20 21

namespace PaddleOCR {

文幕地方's avatar
文幕地方 已提交
22
class DBPostProcessor {
littletomatodonkey's avatar
littletomatodonkey 已提交
23
public:
littletomatodonkey's avatar
littletomatodonkey 已提交
24 25
  void GetContourArea(const std::vector<std::vector<float>> &box,
                      float unclip_ratio, float &distance);
littletomatodonkey's avatar
littletomatodonkey 已提交
26

littletomatodonkey's avatar
littletomatodonkey 已提交
27 28
  cv::RotatedRect UnClip(std::vector<std::vector<float>> box,
                         const float &unclip_ratio);
littletomatodonkey's avatar
littletomatodonkey 已提交
29 30 31 32

  float **Mat2Vec(cv::Mat mat);

  std::vector<std::vector<int>>
littletomatodonkey's avatar
littletomatodonkey 已提交
33
  OrderPointsClockwise(std::vector<std::vector<int>> pts);
littletomatodonkey's avatar
littletomatodonkey 已提交
34

littletomatodonkey's avatar
littletomatodonkey 已提交
35 36
  std::vector<std::vector<float>> GetMiniBoxes(cv::RotatedRect box,
                                               float &ssid);
littletomatodonkey's avatar
littletomatodonkey 已提交
37

littletomatodonkey's avatar
littletomatodonkey 已提交
38
  float BoxScoreFast(std::vector<std::vector<float>> box_array, cv::Mat pred);
39
  float PolygonScoreAcc(std::vector<cv::Point> contour, cv::Mat pred);
littletomatodonkey's avatar
littletomatodonkey 已提交
40 41

  std::vector<std::vector<std::vector<int>>>
littletomatodonkey's avatar
littletomatodonkey 已提交
42
  BoxesFromBitmap(const cv::Mat pred, const cv::Mat bitmap,
43
                  const float &box_thresh, const float &det_db_unclip_ratio,
44
                  const std::string &det_db_score_mode);
littletomatodonkey's avatar
littletomatodonkey 已提交
45 46

  std::vector<std::vector<std::vector<int>>>
littletomatodonkey's avatar
littletomatodonkey 已提交
47 48
  FilterTagDetRes(std::vector<std::vector<std::vector<int>>> boxes,
                  float ratio_h, float ratio_w, cv::Mat srcimg);
littletomatodonkey's avatar
littletomatodonkey 已提交
49 50

private:
littletomatodonkey's avatar
littletomatodonkey 已提交
51 52 53 54 55
  static bool XsortInt(std::vector<int> a, std::vector<int> b);

  static bool XsortFp32(std::vector<float> a, std::vector<float> b);

  std::vector<std::vector<float>> Mat2Vector(cv::Mat mat);
littletomatodonkey's avatar
littletomatodonkey 已提交
56 57 58 59 60 61 62 63 64 65 66 67

  inline int _max(int a, int b) { return a >= b ? a : b; }

  inline int _min(int a, int b) { return a >= b ? b : a; }

  template <class T> inline T clamp(T x, T min, T max) {
    if (x > max)
      return max;
    if (x < min)
      return min;
    return x;
  }
littletomatodonkey's avatar
littletomatodonkey 已提交
68

littletomatodonkey's avatar
littletomatodonkey 已提交
69 70 71 72 73 74 75 76 77
  inline float clampf(float x, float min, float max) {
    if (x > max)
      return max;
    if (x < min)
      return min;
    return x;
  }
};

文幕地方's avatar
文幕地方 已提交
78 79
class TablePostProcessor {
public:
文幕地方's avatar
fix bug  
文幕地方 已提交
80
  void init(std::string label_path, bool merge_no_span_structure = true);
文幕地方's avatar
文幕地方 已提交
81 82 83 84 85 86
  void Run(std::vector<float> &loc_preds, std::vector<float> &structure_probs,
           std::vector<float> &rec_scores, std::vector<int> &loc_preds_shape,
           std::vector<int> &structure_probs_shape,
           std::vector<std::vector<std::string>> &rec_html_tag_batch,
           std::vector<std::vector<std::vector<int>>> &rec_boxes_batch,
           std::vector<int> &width_list, std::vector<int> &height_list);
文幕地方's avatar
文幕地方 已提交
87 88 89 90 91 92 93

private:
  std::vector<std::string> label_list_;
  std::string end = "eos";
  std::string beg = "sos";
};

文幕地方's avatar
文幕地方 已提交
94 95 96 97
class PicodetPostProcessor {
public:
  void init(std::string label_path, const double score_threshold = 0.4,
            const double nms_threshold = 0.5,
文幕地方's avatar
fix bug  
文幕地方 已提交
98
            const std::vector<int> &fpn_stride = {8, 16, 32, 64});
文幕地方's avatar
文幕地方 已提交
99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116
  void Run(std::vector<StructurePredictResult> &results,
           std::vector<std::vector<float>> outs, std::vector<int> ori_shape,
           std::vector<int> resize_shape, int eg_max);
  std::vector<int> fpn_stride_ = {8, 16, 32, 64};

private:
  StructurePredictResult disPred2Bbox(std::vector<float> bbox_pred, int label,
                                      float score, int x, int y, int stride,
                                      std::vector<int> im_shape, int reg_max);
  void nms(std::vector<StructurePredictResult> &input_boxes,
           float nms_threshold);

  std::vector<std::string> label_list_;
  double score_threshold_ = 0.4;
  double nms_threshold_ = 0.5;
  int num_class_ = 5;
};

littletomatodonkey's avatar
littletomatodonkey 已提交
117
} // namespace PaddleOCR