TrainEdgeConcat.h

// Concatenated training model for pairwise potentials
// Written by Sergey G. Kosov in 2015 - 2016 for Project X
#pragma once

#include "TrainEdge.h"
#include "PriorNode.h"
#include "TrainNode.h"
#include "FeaturesConcatenator.h"
#include "macroses.h"

namespace DirectGraphicalModels
{

    // ============================= Concatenated Edge Train Class =============================
    template<class Trainer, class Concatenator> class CTrainEdgeConcat : public CTrainEdge
    {
    public:
        CTrainEdgeConcat(byte nStates, word nFeatures) 
            : CBaseRandomModel(nStates)
            , CTrainEdge(nStates, nFeatures)
        {
            DGM_ASSERT(nStates < 16);
            m_pPrior        = std::make_unique<CPriorNode>(nStates * nStates);
            m_pTrainer      = std::make_unique<Trainer>(nStates * nStates, nFeatures);
            m_pConcatenator = std::make_unique<Concatenator>(nFeatures);
            m_featureVector = Mat(m_pConcatenator->getNumFeatures(), 1, CV_8UC1);
        }
        template<class TrainerParams> CTrainEdgeConcat(byte nStates, word nFeatures, TrainerParams params)
            : CTrainEdge(nStates, nFeatures)
            , CBaseRandomModel(nStates)
        {
            DGM_ASSERT(nStates < 16);
            m_pPrior        = std::make_unique<CPriorNode>(nStates * nStates);
            m_pTrainer      = std::make_unique<Trainer>(nStates * nStates, nFeatures, params);
            m_pConcatenator = std::make_unique<Concatenator>(nFeatures);
            m_featureVector = Mat(m_pConcatenator->getNumFeatures(), 1, CV_8UC1);
        }
        virtual ~CTrainEdgeConcat(void) = default;


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

        virtual void    addFeatureVecs(const Mat &featureVector1, byte gt1, const Mat &featureVector2, byte gt2) 
        {
            byte gt = gt2 * m_nStates + gt1;
            m_pPrior->addNodeGroundTruth(gt);
            m_pConcatenator->concatenate(featureVector1, featureVector2, m_featureVector);
            m_pTrainer->addFeatureVec(m_featureVector, gt);
        }
        virtual void    train(bool doClean = false) { m_pTrainer->train(doClean); }


    protected:
        DllExport virtual void  saveFile(FILE *pFile) const {} 
        DllExport virtual void  loadFile(FILE *pFile) {}        
        DllExport virtual Mat   calculateEdgePotentials(const Mat &featureVector1, const Mat &featureVector2, const vec_float_t &vParams) const 
        {
            const float nodePotWeight = 1.0f;
            m_pConcatenator->concatenate(featureVector1, featureVector2, const_cast<Mat &>(m_featureVector));
            Mat pot = m_pTrainer->getNodePotentials(m_featureVector, nodePotWeight);
            Mat prior = m_pPrior->getPrior(100);

            Mat res(m_nStates, m_nStates, CV_32FC1);
            
            for (byte gt1 = 0; gt1 < m_nStates; gt1++) {
                float * pRes = res.ptr<float>(gt1);
                for (byte gt2 = 0; gt2 < m_nStates; gt2++) {
                    byte gt = gt2 * m_nStates + gt1;
                    
                    float epsilon = prior.at<float>(gt) > 0 ? FLT_EPSILON : 0.0f;
                    pRes[gt2] = MAX(pot.at<float>(gt), epsilon);
                }
            }

            return res;
        }
    

    private:
        std::unique_ptr<CPriorNode>             m_pPrior;           
        std::unique_ptr<CTrainNode>             m_pTrainer;         
        std::unique_ptr<CFeaturesConcatenator>  m_pConcatenator;    
        Mat                                     m_featureVector;    
    };
}