TrainNodeNaiveBayes.h

// Bayes training class interface
// Written by Sergey G. Kosov in 2012 - 2015 for Project X
#pragma once

#include "TrainNode.h"
#include "PriorNode.h"

namespace DirectGraphicalModels
{
    class IPDF;

    // ====================== Bayes Train Class =====================
    class CTrainNodeBayes : public CTrainNode, private CPriorNode
    {
    public:
        DllExport CTrainNodeBayes(byte nStates, word nFeatures);
        DllExport virtual ~CTrainNodeBayes(void);

        DllExport virtual void    reset(void);

        DllExport virtual void    addFeatureVec(const Mat &featureVector, byte gt); 
        DllExport virtual void    train(bool doClean = false);

        DllExport IPDF          * getPDF(byte state, word feature) const { return m_pPDF[state][feature]; } 
        DllExport IPDF          * getPDF2D(byte state) const { return m_pPDF2D[state]; }
        DllExport void            smooth(int nIt = 1);
    
    protected:
        DllExport virtual void    saveFile(FILE *pFile) const; 
        DllExport virtual void    loadFile(FILE *pFile); 
        DllExport void calculateNodePotentials(const Mat &featureVector, Mat &potential, Mat &mask) const;


    private:
        IPDF      *** m_pPDF;               
        IPDF       ** m_pPDF2D;             
        Mat           m_prior;              
    };
}