61 lines
2.4 KiB
C++
61 lines
2.4 KiB
C++
#pragma once
|
|
|
|
#include "ONNXOCRTypes.h"
|
|
#include "ONNXEngine.h"
|
|
#include <vector>
|
|
#include <mutex>
|
|
|
|
namespace ANSCENTER {
|
|
namespace onnxocr {
|
|
|
|
class ONNXOCRDetector : public BasicOrtHandler {
|
|
public:
|
|
explicit ONNXOCRDetector(const std::string& onnx_path, unsigned int num_threads = 1);
|
|
explicit ONNXOCRDetector(const std::string& onnx_path,
|
|
const OrtHandlerOptions& options,
|
|
unsigned int num_threads = 1);
|
|
~ONNXOCRDetector() override = default;
|
|
|
|
// Run text detection on an image
|
|
std::vector<TextBox> Detect(const cv::Mat& srcImage,
|
|
int maxSideLen = kDetMaxSideLen,
|
|
float dbThresh = kDetDbThresh,
|
|
float boxThresh = kDetBoxThresh,
|
|
float unclipRatio = kDetUnclipRatio,
|
|
bool useDilation = false);
|
|
|
|
// Pre-warm cuDNN/TRT at a canonical 320x320 input so the first real
|
|
// call doesn't pay the algorithm-selection tax. Idempotent.
|
|
void Warmup();
|
|
|
|
private:
|
|
bool _warmedUp = false;
|
|
Ort::Value transform(const cv::Mat& mat) override;
|
|
Ort::Value transformBatch(const std::vector<cv::Mat>& images) override;
|
|
|
|
// Preprocessing
|
|
std::vector<float> Preprocess(const cv::Mat& srcImage, int maxSideLen,
|
|
int& resizeH, int& resizeW,
|
|
float& ratioH, float& ratioW);
|
|
|
|
// Postprocessing: threshold -> contours -> boxes (matches PaddleOCR official flow)
|
|
std::vector<TextBox> Postprocess(const float* outputData, int outH, int outW,
|
|
float ratioH, float ratioW, int srcH, int srcW,
|
|
float dbThresh, float boxThresh, float unclipRatio,
|
|
bool useDilation);
|
|
|
|
// Get ordered 4 corners [TL, TR, BR, BL] from rotated rect (matches PaddleOCR GetMiniBoxes)
|
|
std::array<cv::Point2f, 4> GetMiniBoxes(const cv::RotatedRect& rect);
|
|
|
|
// Compute mean score inside box polygon on the probability map
|
|
float BoxScoreFast(const cv::Mat& probMap, const std::array<cv::Point2f, 4>& box);
|
|
|
|
// Expand 4-point box using Clipper offset
|
|
std::vector<cv::Point2f> UnclipPolygon(const std::array<cv::Point2f, 4>& box, float unclipRatio);
|
|
|
|
std::mutex _mutex;
|
|
};
|
|
|
|
} // namespace onnxocr
|
|
} // namespace ANSCENTER
|