ANSCORE/modules/ANSOCR/ANSONNXOCR/PaddleOCRV5Engine.h

#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