ESMH.h

#ifndef MTF_ESMH_H
#define MTF_ESMH_H

#include "SearchMethod.h"

#define ESMH_MAX_ITERS 10
#define ESMH_EPSILON 0.01
#define ESMH_JAC_TYPE 0
#define ESMH_HESS_TYPE 0
#define ESMH_SEC_ORD_HESS false
#define ESMH_ENABLE_SPI false
#define ESMH_SPI_THRESH 10
#define ESMH_DEBUG_MODE false

_MTF_BEGIN_NAMESPACE

struct ESMHParams{

    enum class JacType{ Original, DiffOfJacs };
    enum class HessType {
        Original, SumOfStd, SumOfSelf,
        InitialSelf, CurrentSelf, Std
    };

    int max_iters; 
    double epsilon; 

    JacType jac_type;
    HessType hess_type;
    bool sec_ord_hess;

    bool enable_spi;
    double spi_thresh;
    bool debug_mode; 

    // value constructor
    ESMHParams(int _max_iters, double _epsilon, 
        JacType _jac_type, HessType _hess_type, bool _sec_ord_hess,
        bool _enable_spi, double _spi_thresh,
        bool _debug_mode){
        max_iters = _max_iters;
        epsilon = _epsilon;
        jac_type = _jac_type;
        hess_type = _hess_type;
        sec_ord_hess = _sec_ord_hess;
        enable_spi = _enable_spi;
        spi_thresh = _spi_thresh;
        debug_mode = _debug_mode;
    }
    // default and copy constructor
    ESMHParams(ESMHParams *params = nullptr) :
        max_iters(ESMH_MAX_ITERS), epsilon(ESMH_EPSILON),
        jac_type(static_cast<JacType>(ESMH_JAC_TYPE)),
        hess_type(static_cast<HessType>(ESMH_HESS_TYPE)),
        sec_ord_hess(ESMH_SEC_ORD_HESS),
        enable_spi(ESMH_ENABLE_SPI), spi_thresh(ESMH_SPI_THRESH),
        debug_mode(ESMH_DEBUG_MODE){
        if(params){
            max_iters = params->max_iters;
            epsilon = params->epsilon;
            jac_type = params->jac_type;
            hess_type = params->hess_type;
            sec_ord_hess = params->sec_ord_hess;
            enable_spi = params->enable_spi;
            spi_thresh = params->spi_thresh;
            debug_mode = params->debug_mode;
        }
    }
};

template<class AM, class SSM>
class ESMH : public SearchMethod < AM, SSM > {

public:
    typedef ESMHParams ParamType;

    typedef ParamType::JacType JacType;
    typedef ParamType::HessType HessType;

    using SearchMethod<AM, SSM> ::am;
    using SearchMethod<AM, SSM> ::ssm;
    using typename SearchMethod<AM, SSM> ::AMParams;
    using typename SearchMethod<AM, SSM> ::SSMParams;
    using SearchMethod<AM, SSM> ::cv_corners_mat;
    using SearchMethod<AM, SSM> ::name;
    using SearchMethod<AM, SSM> ::initialize;
    using SearchMethod<AM, SSM> ::update;

protected:
    ParamType params;

    int frame_id;
    VectorXc pix_mask2;
    VectorXb pix_mask;
    VectorXd rel_pix_diff;
    double max_pix_diff;

    Matrix24d prev_corners;
    MatrixXd init_pix_jacobian, curr_pix_jacobian, mean_pix_jacobian;
    MatrixXd init_pix_hessian, curr_pix_hessian, mean_pix_hessian;
    VectorXd ssm_update;
    RowVectorXd jacobian;
    MatrixXd hessian, init_self_hessian;

public:

    ESMH(const ParamType *esm_params,
        const AMParams *am_params, const SSMParams *ssm_params) :
        SearchMethod<AM, SSM>(am_params, ssm_params),
        params(esm_params){
        printf("\n");
        printf("Using ESM tracker with:\n");
        printf("max_iters: %d\n", params.max_iters);
        printf("epsilon: %f\n", params.epsilon);
        printf("jac_type: %d\n", params.jac_type);
        printf("hess_type: %d\n", params.hess_type);
        printf("sec_ord_hess: %d\n", params.sec_ord_hess);
        printf("enable_spi: %d\n", params.enable_spi);
        printf("spi_thresh: %f\n", params.spi_thresh);
        printf("debug_mode: %d\n", params.debug_mode);

        printf("appearance model: %s\n", am.name.c_str());
        printf("state space model: %s\n", ssm.name.c_str());
        printf("\n");

        name = "esm";

        frame_id = 0;
        max_pix_diff = 0;

        switch(params.jac_type){
        case JacType::Original:
            printf("Using original ESM Jacobian\n");
            break;
        case JacType::DiffOfJacs:
            printf("Using Difference of Jacobians\n");
            break;
        default:
            throw utils::InvalidArgument("Invalid Jacobian type provided");
        }

        const char *hess_order = params.sec_ord_hess ? "Second" : "First";
        switch(params.hess_type){
        case HessType::Original:
            printf("Using %s order original ESM Hessian\n", hess_order);
            break;
        case HessType::SumOfStd:
            printf("Using Sum of %s order Standard Hessians\n", hess_order);
            break;
        case HessType::SumOfSelf:
            printf("Using Sum of %s order Self Hessians\n", hess_order);
            break;
        case HessType::InitialSelf:
            printf("Using %s order Initial Self Hessian\n", hess_order);
            break;
        case HessType::CurrentSelf:
            printf("Using %s order Current Self Hessian\n", hess_order);
            break;
        case HessType::Std:
            printf("Using %s order Standard Hessian\n", hess_order);
            break;
        default:
            throw utils::InvalidArgument("Invalid Hessian type provided");
        }

        ssm_update.resize(ssm.getStateSize());
        jacobian.resize(ssm.getStateSize());
        hessian.resize(ssm.getStateSize(), ssm.getStateSize());

        init_pix_jacobian.resize(am.getNPix(), ssm.getStateSize());
        curr_pix_jacobian.resize(am.getNPix(), ssm.getStateSize());

        if(params.jac_type == JacType::Original || params.hess_type == HessType::Original){
            mean_pix_jacobian.resize(am.getNPix(), ssm.getStateSize());
        }
        if(params.hess_type == HessType::SumOfSelf){
            init_self_hessian.resize(ssm.getStateSize(), ssm.getStateSize());
        }
        if(params.sec_ord_hess){
            init_pix_hessian.resize(ssm.getStateSize()*ssm.getStateSize(), am.getNPix());
            if(params.hess_type != HessType::InitialSelf){
                curr_pix_hessian.resize(ssm.getStateSize()*ssm.getStateSize(), am.getNPix());
                if(params.hess_type == HessType::Original){
                    mean_pix_hessian.resize(ssm.getStateSize()*ssm.getStateSize(), am.getNPix());
                }
            }
        }
    }

    void initialize(const cv::Mat &corners) override{
        frame_id = 0;
        ssm.initialize(corners);
        am.initializePixVals(ssm.getPts());

        if(params.enable_spi){
            if(!ssm.supportsSPI())
                throw utils::InvalidArgument("ESM::initialize : SSM does not support SPI");
            if(!am.supportsSPI())
                throw utils::InvalidArgument("ESM::initialize : AM does not support SPI");

            printf("Using Selective Pixel Integration\n");
            pix_mask.resize(am.getNPix());
            ssm.setSPIMask(pix_mask.data());
            am.setSPIMask(pix_mask.data());
        }

        ssm.initializeGradPts(am.getGradOffset());
        am.initializePixGrad(ssm.getGradPts());
        ssm.cmptInitPixJacobian(init_pix_jacobian, am.getInitPixGrad());

        if(params.sec_ord_hess){
            ssm.initializeHessPts(am.getHessOffset());
            am.initializePixHess(ssm.getPts(), ssm.getHessPts());
            ssm.cmptInitPixHessian(init_pix_hessian, am.getInitPixHess(), am.getInitPixGrad());
        }
        am.initializeSimilarity();
        am.initializeGrad();
        am.initializeHess();

        if(params.hess_type == HessType::InitialSelf || params.hess_type == HessType::SumOfSelf){
            if(params.sec_ord_hess){
                am.cmptSelfHessian(hessian, init_pix_jacobian, init_pix_hessian);
            } else{
                am.cmptSelfHessian(hessian, init_pix_jacobian);
            }
            if(params.hess_type == HessType::SumOfSelf){
                init_self_hessian = hessian;
            }
        }
        ssm.getCorners(cv_corners_mat);
    }

    void update() override{
        ++frame_id;
        am.setFirstIter();
        for(int iter_id = 0; iter_id < params.max_iters; iter_id++){
            // extract pixel values from the current image at the latest known position of the object
            am.updatePixVals(ssm.getPts());

            if(params.enable_spi){
                rel_pix_diff = (am.getInitPixVals() - am.getCurrPixVals()) / max_pix_diff;
                pix_mask = rel_pix_diff.cwiseAbs().array() < params.spi_thresh;
            }
            // compute pixel gradient of the current image warped with the current warp
            ssm.updateGradPts(am.getGradOffset());
            am.updatePixGrad(ssm.getGradPts());
            // multiply the pixel gradient with the SSM Jacobian to get the Jacobian of pixel values w.r.t. SSM parameters
            ssm.cmptInitPixJacobian(curr_pix_jacobian, am.getCurrPixGrad());

            if(params.jac_type == JacType::Original || params.hess_type == HessType::Original){
                mean_pix_jacobian = (init_pix_jacobian + curr_pix_jacobian) / 2.0;
            }
            if(params.sec_ord_hess && params.hess_type != HessType::InitialSelf){
                ssm.updateHessPts(am.getHessOffset());
                am.updatePixHess(ssm.getPts(), ssm.getHessPts());
                ssm.cmptInitPixHessian(curr_pix_hessian, am.getCurrPixHess(), am.getCurrPixGrad());
            }
            // compute the prerequisites for the gradient functions
            am.updateSimilarity();
            // update the gradient of the error norm w.r.t. current pixel values
            am.updateCurrGrad();
            // update the gradient of the error norm w.r.t. initial pixel values
            am.updateInitGrad();

            switch(params.jac_type){
            case JacType::Original:
                // take the mean at the level of the jacobian of the pixel values wrt SSM parameters, 
                //then use this mean jacobian to compute the Jacobian of the error norm wrt SSM parameters
                am.cmptCurrJacobian(jacobian, mean_pix_jacobian);
                break;
            case JacType::DiffOfJacs:
                // compute the mean difference between the Jacobians of the error norm w.r.t. initial AND current values of SSM parameters
                am.cmptDifferenceOfJacobians(jacobian, init_pix_jacobian, curr_pix_jacobian);
                jacobian *= 0.5;
                break;
            }
            switch(params.hess_type){
            case HessType::InitialSelf:
                break;
            case HessType::Original:
                if(params.sec_ord_hess){
                    mean_pix_hessian = (init_pix_hessian + curr_pix_hessian) / 2.0;
                    am.cmptCurrHessian(hessian, mean_pix_jacobian, mean_pix_hessian);
                } else{
                    am.cmptCurrHessian(hessian, mean_pix_jacobian);
                }
                break;
            case HessType::SumOfStd:
                if(params.sec_ord_hess){
                    am.cmptSumOfHessians(hessian, init_pix_jacobian, curr_pix_jacobian,
                        init_pix_hessian, curr_pix_hessian);
                } else{
                    am.cmptSumOfHessians(hessian, init_pix_jacobian, curr_pix_jacobian);
                }
                hessian *= 0.5;
                break;
            case HessType::SumOfSelf:
                if(params.sec_ord_hess){
                    am.cmptSelfHessian(hessian, curr_pix_jacobian, curr_pix_hessian);
                } else{
                    am.cmptSelfHessian(hessian, curr_pix_jacobian);
                }
                hessian = (hessian + init_self_hessian) * 0.5;
                break;
            case HessType::CurrentSelf:
                if(params.sec_ord_hess){
                    am.cmptSelfHessian(hessian, curr_pix_jacobian, curr_pix_hessian);
                } else{
                    am.cmptSelfHessian(hessian, curr_pix_jacobian);
                }
                break;
            case HessType::Std:
                if(params.sec_ord_hess){
                    am.cmptCurrHessian(hessian, curr_pix_jacobian, curr_pix_hessian);
                } else{
                    am.cmptCurrHessian(hessian, curr_pix_jacobian);
                }
                break;
            }
            ssm_update = -hessian.colPivHouseholderQr().solve(jacobian.transpose());
            prev_corners = ssm.getCorners();
            ssm.compositionalUpdate(ssm_update);
            double update_norm = (prev_corners - ssm.getCorners()).squaredNorm();
            if(update_norm < params.epsilon){
                break;
            }
            am.clearFirstIter();
        }
        ssm.getCorners(cv_corners_mat);
    }
};
_MTF_END_NAMESPACE

#endif