drwnPatchMatch.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-2017, Stephen Gould
** All rights reserved.
**
******************************************************************************
** FILENAME:    drwnPatchMatch.h
** AUTHOR(S):   Stephen Gould <stephen.gould@anu.edu.au>
**
*****************************************************************************/

#pragma once

#include <cstdlib>
#if defined(_WIN32)||defined(WIN32)||defined(__WIN32__)||defined(__VISUALC__)
#include <cstdint>
#endif
#include <cassert>
#include <list>

#include "cv.h"

#include "drwnBase.h"
#include "drwnIO.h"
#include "drwnVision.h"

using namespace std;

// drwnPatchMatchStatus ------------------------------------------------------

typedef enum { DRWN_PM_DIRTY, 
               DRWN_PM_PASS1, 
               DRWN_PM_PASS2  
} drwnPatchMatchStatus;

// drwnPatchMatchTransform ---------------------------------------------------

typedef unsigned char drwnPatchMatchTransform;

#define DRWN_PM_TRANSFORM_NONE 0x00
#define DRWN_PM_TRANSFORM_HFLIP 0x01
#define DRWN_PM_TRANSFORM_VFLIP 0x02

inline drwnPatchMatchTransform drwnComposeTransforms(const drwnPatchMatchTransform& f,
    const drwnPatchMatchTransform& g) {
    return g ^ f; // xor transformations (horizontal and vertical flips)
}

// drwnPatchMatchNode --------------------------------------------------------

class drwnPatchMatchNode {
 public:
    uint16_t imgIndx;    
    uint8_t imgScale;    
    uint16_t xPosition;  
    uint16_t yPosition;  

 public:
    drwnPatchMatchNode() : imgIndx(-1), imgScale(0), xPosition(-1), yPosition(-1)
        { /* do nothing */ };
    drwnPatchMatchNode(uint16_t idx, uint8_t sc, uint16_t x, uint16_t y) :
        imgIndx(idx), imgScale(sc), xPosition(x), yPosition(y) { /* do nothing */ };
    drwnPatchMatchNode(uint16_t idx, uint8_t sc, const cv::Point& p) :
        imgIndx(idx), imgScale(sc), xPosition(p.x), yPosition(p.y) { /* do nothing */ };
    ~drwnPatchMatchNode() { /* do nothing */ };

    bool operator<(const drwnPatchMatchNode& node) const {
        if (imgIndx < node.imgIndx) return true;
        if (imgIndx > node.imgIndx) return false;
        if (imgScale < node.imgScale) return true;
        if (imgScale > node.imgScale) return false;
        if (yPosition < node.yPosition) return true;
        if (yPosition > node.yPosition) return false;
        return (xPosition < node.xPosition);
    }

    bool operator==(const drwnPatchMatchNode& node) const {
        if (imgIndx != node.imgIndx) return false;
        if (imgScale != node.imgScale) return false;
        if (yPosition != node.yPosition) return false;
        if (xPosition != node.xPosition) return false;
        return true;
    }
};

// drwnPatchMatchEdge --------------------------------------------------------

class drwnPatchMatchEdge {
 public:
    float matchScore;              
    drwnPatchMatchNode targetNode; 
    drwnPatchMatchTransform xform; 

    drwnPatchMatchStatus status;   

 public:
    inline drwnPatchMatchEdge() : matchScore(DRWN_FLT_MAX),
        xform(DRWN_PM_TRANSFORM_NONE), status(DRWN_PM_DIRTY) { /* do nothing */ };
    inline drwnPatchMatchEdge(float score, const drwnPatchMatchNode& node,
        const drwnPatchMatchTransform& patchXform) : matchScore(score), targetNode(node),
        xform(patchXform), status(DRWN_PM_DIRTY) { /* do nothing */ };

    size_t numBytesOnDisk() const;
    bool read(istream& is);
    bool write(ostream& os) const;

    bool operator<(const drwnPatchMatchEdge& e) const {
        return (matchScore < e.matchScore);
    }
};

// drwnPatchMatchEdge Utility Functions --------------------------------------

inline bool drwnPatchMatchSortByScore(const drwnPatchMatchEdge& a, const drwnPatchMatchEdge& b)
{
    return (a.matchScore < b.matchScore);
}

inline bool drwnPatchMatchSortByImage(const drwnPatchMatchEdge& a, const drwnPatchMatchEdge& b)
{
    return (a.targetNode.imgIndx < b.targetNode.imgIndx) ||
        ((a.targetNode.imgIndx == b.targetNode.imgIndx) && (a.matchScore < b.matchScore));
}

string toString(const drwnPatchMatchEdge& e);

// drwnPatchMatchEdgeList ----------------------------------------------------

typedef vector<drwnPatchMatchEdge> drwnPatchMatchEdgeList;

// drwnPatchMatchImageRecord ------------------------------------------------

class drwnPatchMatchImageRecord {
 public:
    static bool ALLOW_MULTIPLE;        

 protected:
    uint16_t  _width;      
    uint16_t  _height;     

    vector<drwnPatchMatchEdgeList> _matches;

 public:
    drwnPatchMatchImageRecord() : _width(0), _height(0) {
        // do nothing
    }
    drwnPatchMatchImageRecord(uint16_t width, uint16_t height) : _width(width), _height(height) {
        DRWN_ASSERT((width > 0) && (height > 0));
        _matches.resize(width * height);
    }
    ~drwnPatchMatchImageRecord() {
        // do nothing
    }

    uint16_t width() const { return _width; }
    uint16_t height() const { return _height; }
    size_t size() const { return _width * _height; }

    void clear();

    inline uint16_t x(int indx) const { return indx % _width; }
    inline uint16_t y(int indx) const { return indx / _width; }
    inline cv::Point index2pixel(int indx) const { return cv::Point(indx % _width, indx / _width); }
    inline int pixel2index(uint16_t x, uint16_t y) const { return y * _width + x; }

    bool update(int indx, const drwnPatchMatchEdge& match);
    inline bool update(uint16_t x, uint16_t y, const drwnPatchMatchEdge& match) {
        return update(pixel2index(x, y), match);
    }

    drwnPatchMatchEdgeList& operator()(uint16_t x, uint16_t y) { return _matches[y * _width + x]; }
    const drwnPatchMatchEdgeList& operator()(uint16_t x, uint16_t y) const { return _matches[y * _width + x]; }
    drwnPatchMatchEdgeList& operator[](int indx) { return _matches[indx]; }
    const drwnPatchMatchEdgeList& operator[](int indx) const { return _matches[indx]; }
};

// drwnPatchMatchImagePyramid -------------------------------------------------

class drwnPatchMatchImagePyramid : public drwnPersistentRecord {
 public:
    static unsigned MAX_LEVELS; 
    static unsigned MAX_SIZE;   
    static unsigned MIN_SIZE;   
    static double PYR_SCALE;    

 public:
    bool bActive;

    int eqvClass;

 protected:
    string _name;      
    uint16_t _width;   
    uint16_t _height;  

    vector<drwnPatchMatchImageRecord> _levels;

 public:
    drwnPatchMatchImagePyramid() : drwnPersistentRecord(),
        bActive(true), eqvClass(-1), _name(""), _width(0), _height(0) {
        // do nothing
    }
    drwnPatchMatchImagePyramid(const string& name, uint16_t width, uint16_t height) :
        drwnPersistentRecord(), bActive(true), eqvClass(-1),
        _name(name), _width(width), _height(height) {
        DRWN_ASSERT((width > 0) && (height > 0));
        constructPyramidLevels();
    }
    ~drwnPatchMatchImagePyramid() {
        // do nothing
    }

    const string& name() const { return _name; }
    uint16_t width() const { return _width; }
    uint16_t height() const { return _height; }
    size_t levels() const { return _levels.size(); }

    void clear();

    // i/o (from drwnPersistentRecord)
    size_t numBytesOnDisk() const;
    bool write(ostream& os) const;
    bool read(istream& is);

    inline cv::Point mapPixel(const cv::Point& p, int srcLevel, int dstLevel) const {
        return cv::Point(p.x * _levels[dstLevel].width() / _levels[srcLevel].width(),
            p.y * _levels[dstLevel].height() / _levels[srcLevel].height());
    }

    bool samplePatchByScore(drwnPatchMatchNode& sample) const;

    drwnPatchMatchImageRecord& operator[](uint8_t indx) { return _levels[indx]; }
    const drwnPatchMatchImageRecord& operator[](uint8_t indx) const { return _levels[indx]; }

 protected:
    void constructPyramidLevels();
};

// drwnPatchMatchGraph -------------------------------------------------------

class drwnPatchMatchGraph {
 public:
    static unsigned int PATCH_WIDTH;   
    static unsigned int PATCH_HEIGHT;  
    static unsigned int K;             

    string imageDirectory;             
    string imageExtension;             

 protected:
    unsigned int _patchWidth;          
    unsigned int _patchHeight;         
    vector<drwnPatchMatchImagePyramid *> _images; 

 public:
    drwnPatchMatchGraph();
    drwnPatchMatchGraph(const drwnPatchMatchGraph& graph);
    ~drwnPatchMatchGraph();

    bool write(const char *filestem) const;
    bool read(const char *filestem);

    size_t size() const { return _images.size(); }

    unsigned int patchWidth() const { return _patchWidth; }
    unsigned int patchHeight() const { return _patchHeight; }

    int findImage(const string& baseName) const;
    void appendImage(const string& baseName);
    void appendImage(const string& baseName, const cv::Size& imgSize);
    void appendImages(const vector<string>& baseNames);
    int removeImage(unsigned imgIndx);

    drwnPatchMatchEdgeList& edges(const drwnPatchMatchNode& node) {
        DRWN_ASSERT(node.xPosition < (*_images[node.imgIndx])[node.imgScale].width());
        DRWN_ASSERT(node.yPosition < (*_images[node.imgIndx])[node.imgScale].height());
        return (*_images[node.imgIndx])[node.imgScale](node.xPosition, node.yPosition);
    }
    const drwnPatchMatchEdgeList& edges(const drwnPatchMatchNode& node) const {
        DRWN_ASSERT(node.xPosition < (*_images[node.imgIndx])[node.imgScale].width());
        DRWN_ASSERT(node.yPosition < (*_images[node.imgIndx])[node.imgScale].height());
        return (*_images[node.imgIndx])[node.imgScale](node.xPosition, node.yPosition);
    }

    pair<double, double> energy() const;

    string imageFilename(int indx) const { return imageFilename(_images[indx]->name()); }
    string imageFilename(const string& baseName) const {
        return imageDirectory.empty() ? baseName + imageExtension :
            imageDirectory + DRWN_DIRSEP + baseName + imageExtension;
    }

    inline drwnPatchMatchImagePyramid& operator[](unsigned indx) { return *_images[indx]; }
    inline const drwnPatchMatchImagePyramid& operator[](unsigned indx) const { return *_images[indx]; }
};

// drwnPatchMatchGraphLearner ------------------------------------------------

class drwnPatchMatchGraphLearner {
 public:
    static double SEARCH_DECAY_RATE;   
    static int FORWARD_ENRICHMENT_K;   
    static bool DO_INVERSE_ENRICHMENT; 
    static bool DO_LOCAL_SEARCH;       
    static bool DO_EXHAUSTIVE;         

    static drwnPatchMatchTransform ALLOWED_TRANSFORMATIONS; 

    static double TOP_VAR_PATCHES;     

#ifdef DRWN_DEBUG_STATISTICS
    static unsigned _dbPatchesScored;
#endif

 protected:
    drwnPatchMatchGraph &_graph;           
    vector<vector<cv::Mat> > _features;    

 public:
    drwnPatchMatchGraphLearner(drwnPatchMatchGraph& graph);
    virtual ~drwnPatchMatchGraphLearner();

    void initialize();
    virtual void initialize(unsigned imgIndx);

    void rescore();

    void update();
    virtual void update(unsigned imgIndx);

 protected:
    // --- feature calculation ---

    void cacheImageFeatures();
    void cacheImageFeatures(unsigned imgIndx);
    virtual void cacheImageFeatures(unsigned imgIndx, const cv::Mat& baseImg);

    int appendCIELabFeatures(const cv::Mat& img, cv::Mat& features, int nChannel = 0) const;
    int appendVerticalFeatures(const cv::Mat& img, cv::Mat& features, int nChannel = 0) const;
    int appendEdgeFeatures(const cv::Mat& img, cv::Mat& features, int nChannel = 0) const;
    int appendTextonFilterFeatures(const cv::Mat& img, cv::Mat& features, int nChannel = 0) const;

    // --- move-making steps ---

    bool propagate(const drwnPatchMatchNode& u, bool bDirection);
    bool search(const drwnPatchMatchNode& u);
    bool local(const drwnPatchMatchNode& u);

    virtual bool enrichment();

    bool exhaustive(const drwnPatchMatchNode& u);

    // --- utility routines ---

    virtual float scoreMatch(const drwnPatchMatchNode &u, const drwnPatchMatchNode& v,
        const drwnPatchMatchTransform& xform, float maxValue = DRWN_FLT_MAX) const;

    cv::Point mapPatch(const cv::Point& p, int imgIndx, int srcScale, int dstScale) const;
};

// drwnPatchMatchGraphRepaint ------------------------------------------------

class drwnPatchMatchGraphRepaint {
 protected:
    const drwnPatchMatchGraph &_graph;    
    vector<cv::Mat> _labels;              

 public:
    drwnPatchMatchGraphRepaint(const drwnPatchMatchGraph& graph);
    virtual ~drwnPatchMatchGraphRepaint();

    virtual cv::Mat retarget(unsigned imgIndx) const;

    void updateImageLabels(unsigned imgIndx, const cv::Mat& labels);

 protected:
    void cacheImageLabels();
    virtual void cacheImageLabels(unsigned imgIndx);
};

// drwnPatchMatchVisualization -----------------------------------------------

namespace drwnPatchMatchVis {
    cv::Mat visualizeMatches(const drwnPatchMatchGraph &graph,
        int imgIndx, const cv::Point& p);

    cv::Mat visualizeMatchQuality(const drwnPatchMatchGraph &graph,
        int imgIndx, float maxScore = 0.0, unsigned kthBest = 0);

    cv::Mat visualizeMatchQuality(const drwnPatchMatchGraph &graph, unsigned kthBest = 0);

    cv::Mat visualizeMatchTransforms(const drwnPatchMatchGraph &graph,
        int imgIndx);

    cv::Mat visualizeMatchTransforms(const drwnPatchMatchGraph &graph);

    cv::Mat visualizeMatchTargets(const drwnPatchMatchGraph &graph,
        int imgIndx);

    cv::Mat visualizeMatchTargets(const drwnPatchMatchGraph &graph);
};