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 <unordered_map>
#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;

        // ----------------------------------------------------------------
        //  Full-frame vs pipeline auto-detection (ported from ANSALPR_OD)
        //
        //  When the caller feeds ANSLPR_OCR pre-cropped vehicle ROIs (each
        //  frame is a different small image), the tracker can't work — the
        //  LP detector sees a totally new image every call so trackIds mean
        //  nothing. In that "pipeline" mode we must skip the ALPRChecker
        //  voting layer entirely and return raw OCR results.
        //
        //  When the caller feeds full-frame video (same resolution every
        //  frame, plates moving through the scene), the tracker works
        //  normally and we run plate text through ALPRChecker majority
        //  voting + spatial dedup to stabilise readings.
        //
        //  Mode is auto-detected by watching whether consecutive frames
        //  share the exact same (width, height) for at least
        //  CONFIRM_THRESHOLD frames. Pipeline crops vary by a few pixels;
        //  full-frame video is pixel-identical.
        // ----------------------------------------------------------------
        struct ImageSizeTracker {
            cv::Size lastSize{ 0, 0 };
            int      consistentCount = 0;
            bool     detectedFullFrame = false;
            static constexpr int CONFIRM_THRESHOLD   = 5;
            static constexpr int MIN_FULLFRAME_WIDTH = 1000;
        };
        std::unordered_map<std::string, ImageSizeTracker> _imageSizeTrackers;

        [[nodiscard]] bool shouldUseALPRChecker(const cv::Size& imageSize,
                                                const std::string& cameraId);

        // ----------------------------------------------------------------
        //  Spatial plate identity persistence (ported from ANSALPR_OD)
        //
        //  Prevents the same plate string from appearing on two different
        //  vehicles in the same frame. The LP tracker may briefly assign
        //  the same trackId to two different plates when vehicles pass
        //  each other, or two different trackIds to the same plate when
        //  occlusion breaks a track. In either case, OCR can produce the
        //  same text for two spatial locations for a frame or two — which
        //  looks like "plate flicker" in the UI.
        //
        //  ensureUniquePlateText() resolves the ambiguity by accumulating
        //  confidence per spatial location. When two detections share a
        //  plate text, the one whose spatial history has the higher score
        //  wins and the other has its className cleared.
        // ----------------------------------------------------------------
        struct SpatialPlateIdentity {
            cv::Point2f center;                   // plate center in frame coords
            std::string plateText;
            float       accumulatedScore    = 0.0f;
            int         framesSinceLastSeen = 0;
        };
        std::mutex _plateIdentitiesMutex;
        std::unordered_map<std::string, std::vector<SpatialPlateIdentity>> _plateIdentities;
        static constexpr float PLATE_SPATIAL_MATCH_THRESHOLD = 0.3f; // IoU threshold

        void ensureUniquePlateText(std::vector<Object>& results,
                                   const std::string& cameraId);

        // --- 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]] std::vector<Object> RunInferencesBatch(
            const cv::Mat& input,
            const std::vector<cv::Rect>& vehicleBoxes,
            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