112 lines
4.3 KiB
C++
112 lines
4.3 KiB
C++
#pragma once
|
||
|
||
#include "ONNXOCRTypes.h"
|
||
#include "ONNXOCRDetector.h"
|
||
#include "ONNXOCRClassifier.h"
|
||
#include "ONNXOCRRecognizer.h"
|
||
|
||
#include <memory>
|
||
#include <mutex>
|
||
#include <atomic>
|
||
#include <chrono>
|
||
#include <thread>
|
||
#include <string>
|
||
#include <vector>
|
||
|
||
namespace ANSCENTER {
|
||
namespace onnxocr {
|
||
|
||
// PaddleOCR V5 pipeline engine: Detection -> (Classification) -> Recognition
|
||
// Mirrors the PaddleOCR::PPOCR interface for drop-in replacement
|
||
class PaddleOCRV5Engine {
|
||
public:
|
||
PaddleOCRV5Engine() = default;
|
||
~PaddleOCRV5Engine() = default;
|
||
|
||
// Initialize the OCR pipeline
|
||
// clsModelPath can be empty to skip classification
|
||
// preferTensorRT: try TensorRT EP first for the three sub-models
|
||
// (cuDNN-friendly cuDNN max-workspace mode either way)
|
||
bool Initialize(const std::string& detModelPath,
|
||
const std::string& clsModelPath,
|
||
const std::string& recModelPath,
|
||
const std::string& dictPath,
|
||
bool preferTensorRT = false);
|
||
|
||
// Run full OCR pipeline on an image
|
||
// Returns results matching PaddleOCR::OCRPredictResult format
|
||
std::vector<OCRPredictResult> ocr(const cv::Mat& img);
|
||
|
||
// Run recognizer only on a pre-cropped text image (no detection step)
|
||
TextLine recognizeOnly(const cv::Mat& croppedImage);
|
||
|
||
// Run recognizer only on a batch of pre-cropped text images in a
|
||
// single batched ORT inference. Skips the detector entirely — the
|
||
// caller is expected to supply crops that are already roughly
|
||
// axis-aligned single-line text (e.g. ALPR plate ROIs, optionally
|
||
// pre-split into rows). Crops are grouped by bucket width, so a
|
||
// single call to this function typically issues 1–2 ORT Runs total.
|
||
std::vector<TextLine> recognizeMany(const std::vector<cv::Mat>& croppedImages);
|
||
|
||
// Configuration setters (matching OCRModelConfig parameters)
|
||
void SetDetMaxSideLen(int val) { _maxSideLen = val; }
|
||
void SetDetDbThresh(float val) { _detDbThresh = val; }
|
||
void SetDetBoxThresh(float val) { _detBoxThresh = val; }
|
||
void SetDetUnclipRatio(float val) { _detUnclipRatio = val; }
|
||
void SetClsThresh(float val) { _clsThresh = val; }
|
||
void SetUseDilation(bool val) { _useDilation = val; }
|
||
|
||
private:
|
||
std::unique_ptr<ONNXOCRDetector> detector_;
|
||
std::unique_ptr<ONNXOCRClassifier> classifier_; // nullptr if not used
|
||
std::unique_ptr<ONNXOCRRecognizer> recognizer_;
|
||
|
||
std::recursive_mutex _mutex;
|
||
std::atomic<bool> _modelLoading{ false };
|
||
|
||
// RAII helper: sets _modelLoading=true on construction, false on destruction.
|
||
struct ModelLoadingGuard {
|
||
std::atomic<bool>& flag;
|
||
explicit ModelLoadingGuard(std::atomic<bool>& f) : flag(f) { flag.store(true); }
|
||
~ModelLoadingGuard() { flag.store(false); }
|
||
ModelLoadingGuard(const ModelLoadingGuard&) = delete;
|
||
ModelLoadingGuard& operator=(const ModelLoadingGuard&) = delete;
|
||
};
|
||
|
||
// Try to lock _mutex with a timeout. Returns a unique_lock that
|
||
// evaluates to true on success.
|
||
std::unique_lock<std::recursive_mutex> TryLockWithTimeout(
|
||
const char* caller, unsigned int timeoutMs = 5000)
|
||
{
|
||
const auto deadline = std::chrono::steady_clock::now()
|
||
+ std::chrono::milliseconds(timeoutMs);
|
||
std::unique_lock<std::recursive_mutex> lk(_mutex, std::defer_lock);
|
||
while (!lk.try_lock()) {
|
||
if (std::chrono::steady_clock::now() >= deadline) {
|
||
std::cerr << "[" << caller << "] Mutex acquisition timed out after "
|
||
<< timeoutMs << " ms"
|
||
<< (_modelLoading.load() ? " (model loading in progress)" : "")
|
||
<< std::endl;
|
||
return lk; // unlocked
|
||
}
|
||
std::this_thread::sleep_for(std::chrono::milliseconds(1));
|
||
}
|
||
return lk; // locked
|
||
}
|
||
|
||
// Detection parameters
|
||
int _maxSideLen = kDetMaxSideLen;
|
||
float _detDbThresh = kDetDbThresh;
|
||
float _detBoxThresh = kDetBoxThresh;
|
||
float _detUnclipRatio = kDetUnclipRatio;
|
||
bool _useDilation = false;
|
||
|
||
// Classifier parameters
|
||
float _clsThresh = kClsThresh;
|
||
|
||
bool _initialized = false;
|
||
};
|
||
|
||
} // namespace onnxocr
|
||
} // namespace ANSCENTER
|