ANSCORE/modules/ANSLPR/ANSLPR_OCR.h

#ifndef ANSLPROCR_H
#define ANSLPROCR_H
#pragma once
#include "ANSLPR.h"
#include <list>
#include <map>
#include <string>
#include <mutex>
#include <utility>
#include <vector>

// Forward-declare ANSONNXOCR to avoid pulling in the full ANSOCR header chain
namespace ANSCENTER { class ANSONNXOCR; struct OCRModelConfig; }

namespace ANSCENTER
{
    /// ANSALPR_OCR — License plate recognition using YOLO for LP detection
    /// and ANSONNXOCR (PaddleOCR v5) for text recognition.
    ///
    /// Pipeline:
    ///   1. Detect license plates using _lpDetector (ANSRTYOLO on NVIDIA GPU, ANSONNXYOLO otherwise)
    ///   2. For each detected plate, run OCR using _ocrEngine (ANSONNXOCR)
    ///   3. Optionally classify plate colour using _lpColourDetector (ANSRTYOLO on NVIDIA GPU, ANSONNXYOLO otherwise)
    ///
    /// Supports multiple countries via the Country enum and ALPR post-processing
    /// from ANSOCR's ANSOCRBase infrastructure.
    class ANSLPR_API ANSALPR_OCR : public ANSALPR {
    private:
        ANSCENTER::EngineType    engineType;

        // --- Detectors ---
        std::unique_ptr<ANSCENTER::ANSODBase> _lpDetector = nullptr;        // License plate detector
        std::unique_ptr<ANSCENTER::ANSODBase> _lpColourDetector = nullptr;  // License plate colour classifier
        std::unique_ptr<ANSCENTER::ANSONNXOCR> _ocrEngine = nullptr;        // OCR text recognizer

        // --- Model configs ---
        ANSCENTER::ModelConfig _lpdmodelConfig;
        ANSCENTER::ModelConfig _lpColourModelConfig;

        std::string             _lpdLabels;
        std::string             _lpColourLabels;
        cv::Mat                 _frameBuffer;  // Reusable buffer for color conversion

        std::vector<std::string> _lprModelClass;

        ALPRChecker     alprChecker;

        // --- Original model zip path (reused for ANSONNXOCR initialization) ---
        std::string     _modelZipFilePath;

        // --- Colour detection helpers ---
        [[nodiscard]] std::string DetectLPColourDetector(const cv::Mat& lprROI, const std::string& cameraId);
        [[nodiscard]] std::string DetectLPColourCached(const cv::Mat& lprROI, const std::string& cameraId, const std::string& plateText);

        // LPC colour cache
        struct ColourCacheEntry {
            std::string colour;
            int         hitCount = 0;
        };
        std::mutex _colourCacheMutex;
        std::unordered_map<std::string, ColourCacheEntry> _colourCache;
        static constexpr size_t COLOUR_CACHE_MAX_SIZE = 200;

        // --- OCR helper ---
        [[nodiscard]] std::string RunOCROnPlate(const cv::Mat& plateROI, const std::string& cameraId);

    public:
        ANSALPR_OCR();
        ~ANSALPR_OCR();
        [[nodiscard]] bool Initialize(const std::string& licenseKey, const std::string& modelZipFilePath, const std::string& modelZipPassword, double detectorThreshold, double ocrThreshold, double colourThreshold) override;
        [[nodiscard]] bool LoadEngine() override;
        [[nodiscard]] bool Inference(const cv::Mat& input, std::string& lprResult) override;
        [[nodiscard]] bool Inference(const cv::Mat& input, std::string& lprResult, const std::string& cameraId) override;
        [[nodiscard]] bool Inference(const cv::Mat& input, const std::vector<cv::Rect>& Bbox, std::string& lprResult) override;
        [[nodiscard]] bool Inference(const cv::Mat& input, const std::vector<cv::Rect>& Bbox, std::string& lprResult, const std::string& cameraId) override;
        [[nodiscard]] std::vector<Object> RunInference(const cv::Mat& input, const std::string& cameraId) override;
        [[nodiscard]] bool Destroy() override;

        /// Propagate country to inner OCR engine so ALPR post-processing
        /// uses the correct plate formats and character corrections.
        void SetCountry(Country country) override;

        /// Propagate debug flag to all sub-detectors
        void ActivateDebugger(bool debugFlag) override {
            _debugFlag = debugFlag;
            if (_lpDetector)       _lpDetector->ActivateDebugger(debugFlag);
            if (_lpColourDetector) _lpColourDetector->ActivateDebugger(debugFlag);
        }
    };
}
#endif