BackgroundSubtractorPAWCS.h

#pragma once

#include "BackgroundSubtractorLBSP.h"
#include "custom_utils/DistanceUtils.h"
#include "custom_utils/RandUtils.h"
#include <iostream>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <iomanip>
#include "ThermalColor.h"
#include "ReliabilityBg.h"

#define BGSPAWCS_DEFAULT_LBSP_REL_SIMILARITY_THRESHOLD (0.333f)
#define BGSPAWCS_DEFAULT_DESC_DIST_THRESHOLD_OFFSET (2)
#define BGSPAWCS_DEFAULT_MIN_COLOR_DIST_THRESHOLD (20)
#define BGSPAWCS_DEFAULT_MAX_NB_WORDS (50)
#define BGSPAWCS_DEFAULT_N_SAMPLES_FOR_MV_AVGS (100)

class BackgroundSubtractorPAWCS : public BackgroundSubtractorLBSP {
public:
    BackgroundSubtractorPAWCS(  float fRelLBSPThreshold=BGSPAWCS_DEFAULT_LBSP_REL_SIMILARITY_THRESHOLD,
                                size_t nDescDistThresholdOffset=BGSPAWCS_DEFAULT_DESC_DIST_THRESHOLD_OFFSET,
                                size_t nMinColorDistThreshold=BGSPAWCS_DEFAULT_MIN_COLOR_DIST_THRESHOLD,
                                size_t nMaxNbWords=BGSPAWCS_DEFAULT_MAX_NB_WORDS,
                                size_t nSamplesForMovingAvgs=BGSPAWCS_DEFAULT_N_SAMPLES_FOR_MV_AVGS);
    virtual ~BackgroundSubtractorPAWCS();
    virtual void initialize(const cv::Mat& oInitImg, const cv::Mat& oROI);
    virtual void refreshModel(size_t nBaseOccCount, float fOccDecrFrac, bool bForceFGUpdate=false);
    virtual void apply(cv::InputArray image, cv::OutputArray fgmask, double learningRateOverride=0);
    virtual void apply(cv::InputArray image, cv::OutputArray fgmask, cv::OutputArray _featureMap, cv::OutputArray _reliabilityMap,
                       cv::OutputArray _foregroundRaw, double learningRateOverride=0);
    virtual void getBackgroundImage(cv::OutputArray backgroundImage) const;
    virtual void getBackgroundDescriptorsImage(cv::OutputArray backgroundDescImage) const;

protected:
    template<size_t nChannels>
    struct ColorLBSPFeature {
        uchar anColor[nChannels];
        ushort anDesc[nChannels];
    };
    struct LocalWordBase {
        size_t nFirstOcc;
        size_t nLastOcc;
        size_t nOccurrences;
    };
    template<typename T>
    struct LocalWord : LocalWordBase {
        T oFeature;
    };
    struct GlobalWordBase {
        float fLatestWeight;
        cv::Mat oSpatioOccMap;
        uchar nDescBITS;
    };
    template<typename T>
    struct GlobalWord : GlobalWordBase {
        T oFeature;
    };
    typedef LocalWord<ColorLBSPFeature<1>> LocalWord_1ch;
    typedef LocalWord<ColorLBSPFeature<3>> LocalWord_3ch;
    typedef GlobalWord<ColorLBSPFeature<1>> GlobalWord_1ch;
    typedef GlobalWord<ColorLBSPFeature<3>> GlobalWord_3ch;
    struct PxInfo_PAWCS : PxInfoBase {
        size_t nGlobalWordMapLookupIdx;
        GlobalWordBase** apGlobalDictSortLUT;
    };
    const size_t m_nMinColorDistThreshold;
    const size_t m_nDescDistThresholdOffset;
    size_t m_nMaxLocalWords, m_nCurrLocalWords;
    size_t m_nMaxGlobalWords, m_nCurrGlobalWords;
    const size_t m_nSamplesForMovingAvgs;
    float m_fLastNonFlatRegionRatio;
    int m_nMedianBlurKernelSize;
    cv::Size m_oDownSampledFrameSize_MotionAnalysis, m_oDownSampledFrameSize_GlobalWordLookup;
    cv::Mat m_oDownSampledROI_MotionAnalysis;
    size_t m_nDownSampledROIPxCount;
    size_t m_nLocalWordWeightOffset;

    LocalWordBase** m_apLocalWordDict;
    LocalWord_1ch* m_aLocalWordList_1ch, *m_pLocalWordListIter_1ch;
    LocalWord_3ch* m_aLocalWordList_3ch, *m_pLocalWordListIter_3ch;
    GlobalWordBase** m_apGlobalWordDict;
    GlobalWord_1ch* m_aGlobalWordList_1ch, *m_pGlobalWordListIter_1ch;
    GlobalWord_3ch* m_aGlobalWordList_3ch, *m_pGlobalWordListIter_3ch;
    PxInfo_PAWCS* m_aPxInfoLUT_PAWCS;

    cv::Mat m_oIllumUpdtRegionMask;
    cv::Mat m_oUpdateRateFrame;
    cv::Mat m_oDistThresholdFrame;
    cv::Mat m_oDistThresholdVariationFrame;
    cv::Mat m_oMeanMinDistFrame_LT, m_oMeanMinDistFrame_ST;
    cv::Mat m_oMeanDownSampledLastDistFrame_LT, m_oMeanDownSampledLastDistFrame_ST;
    cv::Mat m_oMeanRawSegmResFrame_LT, m_oMeanRawSegmResFrame_ST;
    cv::Mat m_oMeanFinalSegmResFrame_LT, m_oMeanFinalSegmResFrame_ST;
    cv::Mat m_oUnstableRegionMask;
    cv::Mat m_oBlinksFrame;
    cv::Mat m_oDownSampledFrame_MotionAnalysis;
    cv::Mat m_oLastRawFGMask;

    cv::Mat m_oFGMask_PreFlood;
    cv::Mat m_oFGMask_FloodedHoles;
    cv::Mat m_oLastFGMask_dilated;
    cv::Mat m_oLastFGMask_dilated_inverted;
    cv::Mat m_oCurrRawFGBlinkMask;
    cv::Mat m_oLastRawFGBlinkMask;
    cv::Mat m_oTempGlobalWordWeightDiffFactor;
    cv::Mat m_oMorphExStructElement;

    cv::Mat featureMap;
    cv::Mat reliabilityMap;
    cv::Mat foregroundRaw;
    cv::Mat reliabilityMap2;
    cv::Mat featureMap3;
    cv::Mat reliabilityMap3;

    ReliabilityBg reliabilityBg;
    float feature1, reliability1, feature2a, reliability2a, feature2b, reliability2b, feature3, reliability3;
    float minScale, maxScale, threshold, scaleFactor;

    void CleanupDictionaries();
    static float GetLocalWordWeight(const LocalWordBase* w, size_t nCurrFrame, size_t nOffset);
    static float GetGlobalWordWeight(const GlobalWordBase* w);
};