TrainNodeKNN.cpp

#include "TrainNodeKNN.h"
#include "KDTree.h"
#include "SamplesAccumulator.h"
#include "mathop.h"

namespace DirectGraphicalModels 
{
    // Constructor
    CTrainNodeKNN::CTrainNodeKNN(byte nStates, word nFeatures, TrainNodeKNNParams params) : CBaseRandomModel(nStates), CTrainNode(nStates, nFeatures)
    {
        init(params);
    }

    // Constructor
    CTrainNodeKNN::CTrainNodeKNN(byte nStates, word nFeatures, size_t maxSamples) : CBaseRandomModel(nStates), CTrainNode(nStates, nFeatures) 
    {
        TrainNodeKNNParams params = TRAIN_NODE_KNN_PARAMS_DEFAULT;
        params.maxSamples = maxSamples;
        init(params);
    }
    
    void CTrainNodeKNN::init(TrainNodeKNNParams params)
    {
        m_pSamplesAcc   = new CSamplesAccumulator(m_nStates, params.maxSamples);
        m_pTree         = new CKDTree();
        m_params        = params;
    }

    // Destructor
    CTrainNodeKNN::~CTrainNodeKNN(void)
    {
        delete m_pSamplesAcc;
        delete m_pTree;
    }

    void CTrainNodeKNN::reset(void)
    {
        m_pSamplesAcc->reset();
        m_pTree->reset();
    }

    void CTrainNodeKNN::save(const std::string &path, const std::string &name, short idx) const
    {
        std::string fileName = generateFileName(path, name.empty() ? "TrainNodeKNN" : name, idx);
        m_pTree->save(fileName);
    }

    void CTrainNodeKNN::load(const std::string &path, const std::string &name, short idx)
    {
        std::string fileName = generateFileName(path, name.empty() ? "TrainNodeKNN" : name, idx);
        m_pTree->load(fileName);
    }

    void CTrainNodeKNN::addFeatureVec(const Mat &featureVector, byte gt)
    {
        m_pSamplesAcc->addSample(featureVector, gt);
    }

    void CTrainNodeKNN::train(bool doClean)
    {
#ifdef DEBUG_PRINT_INFO
        printf("\n");
#endif

        // Filling the <samples> and <classes>
        Mat samples, classes;
        for (byte s = 0; s < m_nStates; s++) {                      // states
            int nSamples = m_pSamplesAcc->getNumSamples(s);
#ifdef DEBUG_PRINT_INFO     
            printf("State[%d] - %d of %d samples\n", s, nSamples, m_pSamplesAcc->getNumInputSamples(s));
#endif
            samples.push_back(m_pSamplesAcc->getSamplesContainer(s));
            classes.push_back(Mat(nSamples, 1, CV_8UC1, Scalar(s)));
            if (doClean) m_pSamplesAcc->release(s);             // free memory
        } // s

        // Training, e.g. building the tree
        m_pTree->build(samples, classes);
    }

    void CTrainNodeKNN::calculateNodePotentials(const Mat &featureVector, Mat &potential, Mat &mask) const
    {
        auto nearestNeighbors = m_pTree->findNearestNeighbors(featureVector.t(), m_params.maxNeighbors);
        //float minr = mathop::Euclidian<byte, float>(featureVector.t(), nearestNeighbors.front()->getKey());

        size_t n = nearestNeighbors.size();
        for (auto node : nearestNeighbors) {
            byte  s = node->getValue();             
            
            //float r = mathop::Euclidian<byte, float>(featureVector.t(), node->getKey());          
            //r = 1 + r - minr;
            //potential.at<float>(s, 0) += 0.1f / (r * r);
            
            potential.at<float>(s, 0) += 1.0f;
        }
        if (n) potential /= static_cast<double>(n);
        potential += m_params.bias;
    }
}