TrainNodeKNN.h

// k-Nearest Neighbors training class interface
// Written by Sergey G. Kosov in 2017 for Project X
#pragma once

#include "TrainNode.h"

namespace DirectGraphicalModels
{
    class CKDTree;
    class CSamplesAccumulator;

    typedef struct TrainNodeKNNParams {
        float   bias;                               
        size_t  maxNeighbors;                       
        size_t  maxSamples;                         

        TrainNodeKNNParams() {}
        TrainNodeKNNParams(float _bias, size_t _maxNeighbors, size_t _maxSamples) : bias(_bias), maxNeighbors(_maxNeighbors), maxSamples(_maxSamples) {}
    } TrainNodeKNNParams;
    
    const TrainNodeKNNParams TRAIN_NODE_KNN_PARAMS_DEFAULT =    TrainNodeKNNParams(
                                                                0.1f,   // Regularization CRF parameter: bias is added to all potential values
                                                                100,    // Max number of neighbors to be used for calculating potentials
                                                                0       // Maximum number of samples to be used in training. 0 means using all the samples
                                                                );

    // ====================== k-Nearest Neighbors Train Class =====================
    class CTrainNodeKNN : public CTrainNode
    {
    public:
        DllExport CTrainNodeKNN(byte nStates, word nFeatures, TrainNodeKNNParams params = TRAIN_NODE_KNN_PARAMS_DEFAULT);
        DllExport CTrainNodeKNN(byte nStates, word nFeatures, size_t maxSamples);
        DllExport ~CTrainNodeKNN(void);

        DllExport void  reset(void);
        DllExport void  save(const std::string &path, const std::string &name = std::string(), short idx = -1) const;
        DllExport void  load(const std::string &path, const std::string &name = std::string(), short idx = -1);

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


    protected:
        DllExport void  saveFile(FILE *pFile) const {}
        DllExport void  loadFile(FILE *pFile) {}
        DllExport void  calculateNodePotentials(const Mat &featureVector, Mat &potential, Mat &mask) const;


    protected:
        CKDTree             * m_pTree;                      
        CSamplesAccumulator * m_pSamplesAcc;                


    private:
        void                  init(TrainNodeKNNParams params);  // This function is called by both constructors


    private:
        TrainNodeKNNParams    m_params;
    };
}