TrainLinkNested.h

// Nested training model for pairwise link potentials
// Written by Sergey G. Kosov in 2016 for Project X
#pragma once

#include "TrainLink.h"

namespace DirectGraphicalModels
{
    class CPriorNode;
    class CTrainNode;
    
    // ============================= Nested Link Train Class =============================
    template<class Trainer> class CTrainLinkNested : public CTrainLink
    {
    public:
        CTrainLinkNested(byte nStatesBase, byte nStatesOccl, word nFeatures) 
            : CBaseRandomModel(nStatesBase * nStatesOccl)
            , CTrainLink(nStatesBase, nStatesOccl, nFeatures)
        {
            word nStates = static_cast<word>(nStatesBase) * static_cast<word>(nStatesOccl);
            DGM_ASSERT(nStates < 256);
            m_pPrior    = new CPriorNode(static_cast<byte>(nStates));
            m_pTrainer  = new Trainer(static_cast<byte>(nStates), nFeatures);           
        }
        template<class TrainerParams> CTrainLinkNested(byte nStatesBase, byte nStatesOccl, word nFeatures, TrainerParams params)
            : CTrainLink(nStatesBase, nStatesOccl, nFeatures)
            , CBaseRandomModel(nStatesBase * nStatesOccl)
        {
            word nStates = static_cast<word>(nStatesBase) * static_cast<word>(nStatesOccl);
            DGM_ASSERT(nStates < 256);
            m_pPrior = new CPriorNode(static_cast<byte>(nStates));
            m_pTrainer = new Trainer(static_cast<byte>(nStates), nFeatures, params);
        }
        
        virtual ~CTrainLinkNested(void) 
        { 
            delete m_pPrior;
            delete m_pTrainer; 
        }

        virtual void    reset(void) { m_pPrior->reset(); m_pTrainer->reset(); }
        virtual void    save(const std::string &path, const std::string &name = std::string(), short idx = -1) const { m_pTrainer->save(path, name.empty() ? "CTrainLinkNested" : name, idx); }
        virtual void    load(const std::string &path, const std::string &name = std::string(), short idx = -1) { m_pTrainer->load(path, name.empty() ? "CTrainLinkNested" : name, idx); }

        virtual void    addFeatureVec(const Mat &featureVector, byte gtb, byte gto)
        {
            byte gt = gtb + m_nStatesBase * gto;
            m_pPrior->addNodeGroundTruth(gt);
            m_pTrainer->addFeatureVec(featureVector, gt);
        }
        
        virtual void    train(bool doClean = false)
        { 
            // Fill holes in trainig
            Mat fv(getNumFeatures(), 1, CV_8UC1, Scalar(0));
            Mat priors = m_pPrior->getPrior();
            for (byte i = 0; i < priors.rows; i++)
                if (priors.at<float>(i, 0) == 0)
                    m_pTrainer->addFeatureVec(fv, i);
                
            
            m_pTrainer->train(doClean); 
        }


    protected:
        DllExport virtual void  saveFile(FILE *pFile) const {}
        DllExport virtual void  loadFile(FILE *pFile) {}
        virtual Mat calculateLinkPotentials(const Mat &featureVector) const
        {
            Mat pot = m_pTrainer->getNodePotentials(featureVector);
            //pot = m_pPrior->getPrior(100);

            DGM_ASSERT_MSG(pot.rows == m_nStatesBase * m_nStatesOccl, "The length of the node potentinal vector = %d, but must be %d", pot.rows, m_nStatesBase * m_nStatesOccl);

            Mat res(m_nStatesBase + m_nStatesOccl, m_nStatesBase + m_nStatesOccl, CV_32FC1, Scalar(0));

            for (byte gto = 0; gto < m_nStatesOccl; gto++) {
                float * pRes = res.ptr<float>(m_nStatesBase + gto);
                for (byte gtb = 0; gtb < m_nStatesBase; gtb++) {
                    byte gt = gtb + m_nStatesBase * gto;
                    pRes[gtb] = pot.at<float>(gt, 0);
                }
            }

            return res;
        }

    //  CPriorNode * getPrior(void) const { return m_pPrior; }
    //  CTrainNode * getTrainer(void) const { return m_pTrainer; }

    private:
        CPriorNode              * m_pPrior;             
        CTrainNode              * m_pTrainer;           
    };
}