ANSCustomModels/ANSCustomFireNSmokeDetection/ANSCustomFireNSmoke.h

#include "ANSLIB.h"
#include <deque>
#include <unordered_set>
#include <unordered_map>

#define RETAINFRAMES 80
#define FILTERFRAMES 10
//#define FNS_DEBUG
class CUSTOM_API ANSCustomFS : public IANSCustomClass
{
    typedef std::pair<cv::Scalar, cv::Scalar> Range;
    struct ImageSection {
        cv::Rect region;
        int priority;
        ImageSection(const cv::Rect& r) : region(r), priority(0) {}
    };

    // Track record for voting-based detection confirmation
    struct TrackRecord {
        int trackId{ 0 };
        int classId{ 0 };                        // fire=0, smoke=2
        std::deque<cv::Rect> bboxHistory;         // bounding box history within window
        int detectedCount{ 0 };                   // frames detected in sliding window
        int totalFrames{ 0 };                     // total frames since track appeared
        bool confirmed{ false };                  // passed voting threshold
    };

    // Voting mechanism constants
    static constexpr int VOTE_WINDOW = 15;
    static constexpr int VOTE_THRESHOLD = 8;
    static constexpr float BBOX_CHANGE_THRESHOLD = 0.05f;

private:
    using ANSLIBPtr = std::unique_ptr<ANSCENTER::ANSLIB, decltype(&ANSCENTER::ANSLIB::Destroy)>;

    int    engineType{ 0 };
    ANSLIBPtr _motiondetector{ nullptr, &ANSCENTER::ANSLIB::Destroy };
    ANSLIBPtr _detector{ nullptr, &ANSCENTER::ANSLIB::Destroy };
    ANSLIBPtr _filter{ nullptr, &ANSCENTER::ANSLIB::Destroy };

    std::recursive_mutex    _mutex;
    cv::Rect                _detectedArea;
    cv::Rect                _previousDetectedArea;
    int                     _retainDetectedArea{ 0 };
    bool                    _isFireNSmokeDetected{ false };
    float                   _detectionScoreThreshold{ 0.5 };
    std::vector<cv::Mat>    _frameHistory;
    std::deque <cv::Mat>    _frameQueue;
    std::string             _filterLabels;
    int                     _realFireCheck{ 0 };
    bool                    _isRealFireFrame{ false };
    double                  _hsvThreshold;
    std::vector<Range>      _smoke_colour;
    std::vector<Range>      _fire_colour;

    // ----- Model config -----
    int _detectorModelType;
    int _detectorDetectionType;
    int _filterModelType;
    int _filterDetectionType;
    int _motionModelType;
    int _motionDetectionType;

    // Add these constants at class level
    int MIN_ROI_SIZE = 50;
    float SMOKE_CONFIDENCE_THRESHOLD = 0.7f;
    float DETECTION_CONFIDENCE_THRESHOLD = 0.5f;
    float MAX_AREA_OVERLAP = 0.5f;
    int MAX_MOTION_TRACKING = 10;
    int FILTER_VERIFICATION_FRAMES = 10;
    int MIN_IMAGE_SIZE = 10;
    float INCLUSIVE_IOU_THRESHOLD = 0.001f;
    float MOTION_SENSITIVITY = 0.5f;

    // Excluded filter class IDs
    const std::unordered_set<int> EXCLUDED_FILTER_CLASSES =
    {
        13,  // bench
        59,  // bed
        60,  // dining table
        68,  // microwave
        69,  // oven
        70   // toaster
    };

    cv::Size previousImageSize = cv::Size(0, 0);
    std::vector<ImageSection> cachedSections;
    int                     _currentPriority{ 0 };
    bool                    _readROIs{ false };
    std::vector<cv::Rect>   _exclusiveROIs;
    float                   _smokeDetetectionThreshold{ 0 };
    float                   _motionSpecificity{ 0 };

    // Tracker-based voting state
    std::unordered_map<int, TrackRecord> _trackHistory;

    cv::Rect                GenerateMinimumSquareBoundingBox(const std::vector<ANSCENTER::Object>& detectedObjects, int minSize = 640);
    void                    UpdateNoDetectionCondition();
    bool                    detectStaticFire(std::deque<cv::Mat>& frameQueue);
    bool                    detectScreenFlicker(const std::vector<cv::Mat>& frames);
    bool                    detectReflection(const cv::Mat& frame);
    bool                    MajorityColourInFrame(cv::Mat frame, Range range, float area_threshold);
    bool                    MajorityColourInFrame(cv::Mat frame, std::vector<Range> ranges, float area_threshold);
    bool                    DetectFireNSmokeColourInFrame(const cv::Mat& frame, const cv::Rect& bBox, const std::vector<Range>& ranges, float area_threshold);
    bool                    IsFireDetected(const cv::Mat image, const cv::Rect bBox);
    bool                    IsSmokeDetected(const cv::Mat image, const cv::Rect bBox);
    float                   calculateIoU(const cv::Rect& box1, const cv::Rect& box2);
    bool                    IsOverlapping(const ANSCENTER::Object& obj, const std::vector<ANSCENTER::Object>& objectList, float iouThreshold = 0.5);
    bool                    IsOverlapping(const ANSCENTER::Object& obj, const std::vector<cv::Rect>& objectList, float iouThreshold = 0.5);

    bool                    IsROIOverlapping(const cv::Rect& obj, const std::vector<cv::Rect>& objectList, float iouThreshold = 0.5);
    bool                    IsDetectedAreaValid(const cv::Rect& area) const;
    bool                    IsFireOrSmoke(int classId, float confidence) const;
    bool                    ValidateMotionCorrelation(const std::vector<cv::Rect>& fireNSmokeRects) const;
    void                    UpdatePositiveDetection();
    void                    ResetDetectedArea();
    void                    UpdatePriorityRegion(const cv::Mat& frame);
    int                     CalculateOptimalCropSize(const cv::Mat& frame) const;
    void                    RefineDetectedArea(const cv::Mat& frame, const ANSCENTER::Object& detection);
    std::vector<CustomObject> ConvertToCustomObjects(const std::vector<ANSCENTER::Object>& objects);
    void                    ResetDetectionState();
    void                    GetModelParameters();
    std::vector<ANSCENTER::Object> ProcessExistingDetectedArea(const cv::Mat& frame, const std::string& camera_id, cv::Mat& draw);
    void AddConfirmedDetection(ANSCENTER::Object& detectedObj, std::vector<ANSCENTER::Object>& output);
#ifdef FNS_DEBUG
    void DisplayDebugFrame(cv::Mat& draw) {
        cv::resize(draw, draw, cv::Size(1920, (1920 * draw.rows) / draw.cols));
        cv::imshow("Combined Detected Areas", draw);
        cv::waitKey(1);
    }
#endif
    // Function to separate screen size
    double calculateDistanceToCenter(const cv::Point& center, const cv::Rect& rect);
    std::vector<ImageSection> divideImage(const cv::Mat& image);
    int getHighestPriorityRegion();
    int getLowestPriorityRegion();
    cv::Rect getRegionByPriority(int priority);

    // Tracker-based voting methods
    void UpdateTrackHistory(int trackId, int classId, const cv::Rect& bbox);
    bool IsTrackConfirmed(int trackId) const;
    bool HasBboxMovement(int trackId) const;
    void AgeTracks(const std::unordered_set<int>& detectedTrackIds);
    std::vector<ANSCENTER::Object> FindNewDetectedArea(
        const cv::Mat& frame,
        const std::string& camera_id, cv::Mat& draw);

    std::vector<ANSCENTER::Object> FindMovementObjects(
        const cv::Mat& frame,
        const std::string& camera_id, cv::Mat& draw);

    std::vector<ANSCENTER::Object> FindExcludedObjects(
        const cv::Mat& frame,
        const std::string& camera_id, cv::Mat& draw);

public:
    bool Initialize(const std::string& modelDiretory, float detectionScoreThreshold, std::string& labelMap) override;
    bool OptimizeModel(bool fp16) override;
    bool ConfigureParameters(CustomParams& param) override;
    std::vector<CustomObject> RunInference(const cv::Mat& input) override;
    std::vector<CustomObject> RunInference(const cv::Mat& input, const std::string& camera_id) override;
    bool Destroy() override;
    ANSCustomFS();
    ~ANSCustomFS();
};