145 lines
5.0 KiB
C++
145 lines
5.0 KiB
C++
#include "PaddleOCRV5Engine.h"
|
|
#include "EPLoader.h"
|
|
|
|
#include <opencv2/imgproc.hpp>
|
|
#include <iostream>
|
|
#include <algorithm>
|
|
|
|
namespace ANSCENTER {
|
|
namespace onnxocr {
|
|
|
|
bool PaddleOCRV5Engine::Initialize(const std::string& detModelPath,
|
|
const std::string& clsModelPath,
|
|
const std::string& recModelPath,
|
|
const std::string& dictPath) {
|
|
std::lock_guard<std::recursive_mutex> lock(_mutex);
|
|
ModelLoadingGuard mlg(_modelLoading);
|
|
|
|
try {
|
|
// Initialize detector (also triggers EPLoader init in BasicOrtHandler)
|
|
detector_ = std::make_unique<ONNXOCRDetector>(detModelPath);
|
|
std::cout << "[PaddleOCRV5Engine] Detector initialized: " << detModelPath << std::endl;
|
|
|
|
// Ensure this DLL's copy of Ort::Global<void>::api_ is initialized.
|
|
// BasicOrtHandler sets it in ONNXEngine.dll, but each DLL has its own
|
|
// inline-static copy. Without this, inference calls from ANSOCR.dll crash.
|
|
if (Ort::Global<void>::api_ == nullptr) {
|
|
Ort::InitApi(static_cast<const OrtApi*>(EPLoader::GetOrtApiRaw()));
|
|
}
|
|
|
|
// Initialize classifier (optional)
|
|
if (!clsModelPath.empty()) {
|
|
classifier_ = std::make_unique<ONNXOCRClassifier>(clsModelPath);
|
|
std::cout << "[PaddleOCRV5Engine] Classifier initialized: " << clsModelPath << std::endl;
|
|
}
|
|
else {
|
|
classifier_.reset();
|
|
std::cout << "[PaddleOCRV5Engine] Classifier skipped (no model path)" << std::endl;
|
|
}
|
|
|
|
// Initialize recognizer
|
|
recognizer_ = std::make_unique<ONNXOCRRecognizer>(recModelPath);
|
|
if (!recognizer_->LoadDictionary(dictPath)) {
|
|
std::cerr << "[PaddleOCRV5Engine] Failed to load dictionary" << std::endl;
|
|
return false;
|
|
}
|
|
std::cout << "[PaddleOCRV5Engine] Recognizer initialized: " << recModelPath << std::endl;
|
|
|
|
_initialized = true;
|
|
std::cout << "[PaddleOCRV5Engine] Pipeline initialized successfully" << std::endl;
|
|
return true;
|
|
}
|
|
catch (const std::exception& e) {
|
|
std::cerr << "[PaddleOCRV5Engine] Initialization failed: " << e.what() << std::endl;
|
|
detector_.reset();
|
|
classifier_.reset();
|
|
recognizer_.reset();
|
|
_initialized = false;
|
|
return false;
|
|
}
|
|
}
|
|
|
|
std::vector<OCRPredictResult> PaddleOCRV5Engine::ocr(const cv::Mat& img) {
|
|
if (_modelLoading.load()) return {};
|
|
|
|
std::vector<OCRPredictResult> results;
|
|
|
|
{
|
|
auto lk = TryLockWithTimeout("PaddleOCRV5Engine::ocr");
|
|
if (!lk.owns_lock()) return results;
|
|
if (!_initialized || img.empty()) return results;
|
|
}
|
|
// _mutex released — heavy pipeline runs lock-free
|
|
|
|
// Step 1: Text Detection
|
|
auto boxes = detector_->Detect(img, _maxSideLen, _detDbThresh, _detBoxThresh, _detUnclipRatio, _useDilation);
|
|
|
|
if (boxes.empty()) {
|
|
return results;
|
|
}
|
|
|
|
// Step 2: Crop detected text regions
|
|
std::vector<cv::Mat> croppedImages;
|
|
croppedImages.reserve(boxes.size());
|
|
for (auto& box : boxes) {
|
|
cv::Mat cropped = GetRotateCropImage(img, box);
|
|
if (!cropped.empty()) {
|
|
croppedImages.push_back(cropped);
|
|
}
|
|
}
|
|
|
|
// Step 3: Classification (optional)
|
|
std::vector<int> cls_labels(croppedImages.size(), 0);
|
|
std::vector<float> cls_scores(croppedImages.size(), 0.0f);
|
|
|
|
if (classifier_) {
|
|
classifier_->Classify(croppedImages, cls_labels, cls_scores, _clsThresh);
|
|
|
|
// Rotate images classified as upside-down (label=1 and score > threshold)
|
|
for (size_t i = 0; i < croppedImages.size(); i++) {
|
|
if (cls_labels[i] % 2 == 1 && cls_scores[i] > _clsThresh) {
|
|
cv::rotate(croppedImages[i], croppedImages[i], cv::ROTATE_180);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Step 4: Text Recognition
|
|
auto textLines = recognizer_->RecognizeBatch(croppedImages);
|
|
|
|
// Step 5: Combine results
|
|
for (size_t i = 0; i < boxes.size() && i < textLines.size(); i++) {
|
|
OCRPredictResult result;
|
|
|
|
// Convert TextBox points to box format [[x0,y0], [x1,y1], [x2,y2], [x3,y3]]
|
|
result.box.resize(4);
|
|
for (int j = 0; j < 4; j++) {
|
|
result.box[j] = {
|
|
static_cast<int>(boxes[i].points[j].x),
|
|
static_cast<int>(boxes[i].points[j].y)
|
|
};
|
|
}
|
|
|
|
result.text = textLines[i].text;
|
|
result.score = textLines[i].score;
|
|
result.cls_label = cls_labels[i];
|
|
result.cls_score = cls_scores[i];
|
|
|
|
results.push_back(result);
|
|
}
|
|
|
|
return results;
|
|
}
|
|
|
|
TextLine PaddleOCRV5Engine::recognizeOnly(const cv::Mat& croppedImage) {
|
|
if (_modelLoading.load()) return { "", 0.0f };
|
|
{
|
|
auto lk = TryLockWithTimeout("PaddleOCRV5Engine::recognizeOnly");
|
|
if (!lk.owns_lock()) return { "", 0.0f };
|
|
if (!_initialized || !recognizer_ || croppedImage.empty()) return { "", 0.0f };
|
|
}
|
|
return recognizer_->Recognize(croppedImage);
|
|
}
|
|
|
|
} // namespace onnxocr
|
|
} // namespace ANSCENTER
|