drwnMaskedPatchMatch.h

/*****************************************************************************
** DARWIN: A FRAMEWORK FOR MACHINE LEARNING RESEARCH AND DEVELOPMENT
** Distributed under the terms of the BSD license (see the LICENSE file)
** Copyright (c) 2007-2015, Stephen Gould
** All rights reserved.
**
******************************************************************************
** FILENAME:    drwnMaskedPatchMatch.h
** AUTHOR(S):   Stephen Gould <stephen.gould@anu.edu.au>
**
*****************************************************************************/

// c++ standard headers
#include <cstdlib>
#include <cstdio>

// eigen matrix library headers
#include "Eigen/Core"

// opencv library headers
#include "cv.h"

// darwin library headers
#include "drwnBase.h"
#include "drwnIO.h"
#include "drwnVision.h"

using namespace std;

// drwnBasicPatchMatch -------------------------------------------------------

cv::Mat drwnBasicPatchMatch(const cv::Mat& imgA, const cv::Mat& imgB,
    const cv::Size& patchRadius, cv::Mat& nnfA, unsigned maxIterations = 2);

// drwnSelfPatchMatch --------------------------------------------------------

cv::Mat drwnSelfPatchMatch(const cv::Mat& imgA, const cv::Size& patchRadius,
    cv::Mat& nnfA, double illegalOverlap = 0.0, unsigned maxIterations = 2);

// drwnNNFRepaint ------------------------------------------------------------

cv::Mat drwnNNFRepaint(const cv::Mat& img, const cv::Mat& nnf);

// drwnMaskedPatchMatch ------------------------------------------------------

class drwnMaskedPatchMatch {
 public:
    static bool TRY_IDENTITY_INIT; 
    static int DISTANCE_MEASURE;   
    static float HEIGHT_PENALTY;   

 protected:
    cv::Mat _imgA;          
    cv::Mat _imgB;          
    cv::Mat _maskA;         
    cv::Mat _maskB;         
    cv::Mat _invmaskA;      
    cv::Mat _invmaskB;      

    cv::Mat _overlapA;      
    cv::Mat _validB;        

    cv::Size _patchRadius;  
    cv::Mat _nnfA;          
    cv::Mat _costsA;        

    int _iterationCount;    
    cv::Mat _lastChanged;   

 public:
    drwnMaskedPatchMatch(const cv::Mat& imgA, const cv::Mat& imgB, unsigned patchRadius);
    drwnMaskedPatchMatch(const cv::Mat& imgA, const cv::Mat& imgB, const cv::Size& patchRadius);
    drwnMaskedPatchMatch(const cv::Mat& imgA, const cv::Mat& imgB, const cv::Mat& maskA,
        const cv::Mat& maskB, unsigned patchRadius);
    drwnMaskedPatchMatch(const cv::Mat& imgA, const cv::Mat& imgB, const cv::Mat& maskA,
        const cv::Mat& maskB, const cv::Size& patchRadius);
    virtual ~drwnMaskedPatchMatch() { /* do nothing */ }

    const cv::Mat& getSourceImage() const { return _imgA; }
    const cv::Mat& getTargetImage() const { return _imgB; }
    const cv::Mat& getSourceMask() const { return _maskA; }
    const cv::Mat& getTargetMask() const { return _maskB; }
    cv::Size getPatchSize() const { return cv::Size(2 * _patchRadius.width + 1, 2 * _patchRadius.height + 1); }
    cv::Size getFieldSize() const { return _nnfA.size(); }
    std::pair<cv::Rect, cv::Rect> getMatchingPatches(const cv::Point& ptA) const;

    void initialize() { initialize(_patchRadius); }
    void initialize(const cv::Size& patchRadius);
    void initialize(const cv::Mat& nnf);

    const cv::Mat& search(unsigned maxIterations = 1) {
        return search(cv::Rect(0, 0, _nnfA.cols, _nnfA.rows), maxIterations);
    }
    const cv::Mat& search(cv::Rect roiToUpdate, unsigned maxIterations = 1);

    const cv::Mat& nnf() const { return _nnfA; }
    const cv::Mat& costs() const { return _costsA; }
    double energy() const { return cv::norm(_costsA, cv::NORM_L1, _maskA); }

    void modifySourceImage(const cv::Rect& roi, const cv::Mat& img, double alpha = 0.0);
    void modifySourceImage(const cv::Rect& roiA, const cv::Rect& roiB, double alpha = 0.0);
    void modifyTargetImage(const cv::Rect& roi, const cv::Mat& img, double alpha = 0.0);
    void modifyTargetImage(const cv::Rect& roiA, const cv::Rect& roiB, double alpha = 0.0);
    void expandTargetMask(unsigned radius = 1);


    cv::Mat visualize() const;

 protected:
    bool update(const cv::Point& ptA, const cv::Point& ptB);
    float score(const cv::Point& ptA, const cv::Point& ptB) const;

    cv::Rect affectedRegion(const cv::Rect& modified) const {
        const int y1 = std::max(modified.y - _patchRadius.height, _patchRadius.height);
        const int x1 = std::max(modified.x - _patchRadius.width, _patchRadius.width);
        const int y2 = std::min(modified.y + modified.height + _patchRadius.height, _nnfA.rows - _patchRadius.height);
        const int x2 = std::min(modified.x + modified.width + _patchRadius.width, _nnfA.cols - _patchRadius.width);
        return cv::Rect(x1, y1, x2 - x1, y2 - y1);
    }

    void rescore() { rescore(cv::Rect(0, 0, _imgA.rows, _imgA.cols)); }
    void rescore(const cv::Rect& roi);

    void cacheValidPixels();
    void updateValidPixels(const cv::Rect& roi);
};