102 lines
3.6 KiB
C++
102 lines
3.6 KiB
C++
#pragma once
|
|
|
|
#include "RTOCRTypes.h"
|
|
#include "RTOCRDetector.h"
|
|
#include "RTOCRClassifier.h"
|
|
#include "RTOCRRecognizer.h"
|
|
#include <memory>
|
|
#include <mutex>
|
|
#include <atomic>
|
|
#include <chrono>
|
|
#include <thread>
|
|
#include <string>
|
|
#include <vector>
|
|
#include "ANSLicense.h"
|
|
namespace ANSCENTER {
|
|
namespace rtocr {
|
|
|
|
class PaddleOCRV5RTEngine {
|
|
public:
|
|
PaddleOCRV5RTEngine() = default;
|
|
~PaddleOCRV5RTEngine() = default;
|
|
PaddleOCRV5RTEngine(const PaddleOCRV5RTEngine&) = delete;
|
|
PaddleOCRV5RTEngine& operator=(const PaddleOCRV5RTEngine&) = delete;
|
|
|
|
// Initialize all components
|
|
// clsModelPath can be empty to skip classifier
|
|
bool Initialize(const std::string& detModelPath,
|
|
const std::string& clsModelPath,
|
|
const std::string& recModelPath,
|
|
const std::string& dictPath,
|
|
int gpuId = 0,
|
|
const std::string& engineCacheDir = "");
|
|
|
|
// Run full OCR pipeline: detect → crop → [classify →] recognize
|
|
std::vector<OCRPredictResult> ocr(const cv::Mat& image);
|
|
|
|
// Run recognizer only on a pre-cropped text image (no detection step)
|
|
TextLine recognizeOnly(const cv::Mat& croppedImage);
|
|
|
|
// Configuration setters
|
|
void SetDetMaxSideLen(int v) { detMaxSideLen_ = v; }
|
|
void SetDetDbThresh(float v) { detDbThresh_ = v; }
|
|
void SetDetBoxThresh(float v) { detBoxThresh_ = v; }
|
|
void SetDetUnclipRatio(float v) { detUnclipRatio_ = v; }
|
|
void SetClsThresh(float v) { clsThresh_ = v; }
|
|
void SetUseDilation(bool v) { useDilation_ = v; }
|
|
void SetRecImageHeight(int v) { recImgH_ = v; }
|
|
void SetRecImageMaxWidth(int v) { recImgMaxW_ = v; }
|
|
void SetGpuId(int v) { gpuId_ = v; }
|
|
void SetEngineCacheDir(const std::string& v) { engineCacheDir_ = v; }
|
|
|
|
private:
|
|
std::unique_ptr<RTOCRDetector> detector_;
|
|
std::unique_ptr<RTOCRClassifier> classifier_; // optional
|
|
std::unique_ptr<RTOCRRecognizer> recognizer_;
|
|
|
|
// Configuration
|
|
int detMaxSideLen_ = kDetMaxSideLen;
|
|
float detDbThresh_ = kDetDbThresh;
|
|
float detBoxThresh_ = kDetBoxThresh;
|
|
float detUnclipRatio_ = kDetUnclipRatio;
|
|
float clsThresh_ = kClsThresh;
|
|
bool useDilation_ = false;
|
|
int recImgH_ = kRecImgH;
|
|
int recImgMaxW_ = kRecImgMaxW;
|
|
int gpuId_ = 0;
|
|
std::string engineCacheDir_;
|
|
|
|
std::recursive_mutex _mutex;
|
|
std::atomic<bool> _modelLoading{ false };
|
|
|
|
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;
|
|
};
|
|
|
|
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;
|
|
}
|
|
std::this_thread::sleep_for(std::chrono::milliseconds(1));
|
|
}
|
|
return lk;
|
|
}
|
|
};
|
|
|
|
} // namespace rtocr
|
|
} // namespace ANSCENTER
|