drwnNNGraph.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:    drwnNNGraph.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 <stdint.h>
#include <cassert>
#include <list>

#include "Eigen/Core"

#include "cv.h"

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

using namespace std;
using namespace Eigen;

// drwnNNGraphEdgeStatus -----------------------------------------------------

typedef enum { DRWN_NNG_DIRTY,           
               DRWN_NNG_PROCESSED_ONCE,  
               DRWN_NNG_PROCESSED_TWICE  
} drwnNNGraphEdgeStatus;

// drwnNNGraphNodeIndex ------------------------------------------------------

class drwnNNGraphNodeIndex {
 public:
    uint16_t imgIndx;                    
    uint16_t segId;                      

 public:
    drwnNNGraphNodeIndex() : imgIndx(0), segId(0) { /* do nothing */ }
    drwnNNGraphNodeIndex(uint16_t i, uint16_t s) : imgIndx(i), segId(s) { /* do nothing */ }

    bool operator<(const drwnNNGraphNodeIndex& n) const {
        if (imgIndx < n.imgIndx) return true;
        if (imgIndx > n.imgIndx) return false;
        return (segId < n.segId);
    }

    bool operator==(const drwnNNGraphNodeIndex& n) const {
        return ((imgIndx == n.imgIndx) && (segId == n.segId));
    }

    bool operator!=(const drwnNNGraphNodeIndex& n) const {
        return ((imgIndx != n.imgIndx) || (segId != n.segId));
    }

    friend ostream& operator<<(ostream& os, const drwnNNGraphNodeIndex& n) { 
        os << "(" << n.imgIndx << ", " << n.segId << ")"; return os;
    }
};

// drwnNNGraphEdge -----------------------------------------------------------

class drwnNNGraphEdge {
 public:
    drwnNNGraphNodeIndex targetNode;     
    float weight;                        

    drwnNNGraphEdgeStatus status;        

 public:
    drwnNNGraphEdge() :
        weight(DRWN_FLT_MAX), status(DRWN_NNG_DIRTY) { /* do nothing */ }
    drwnNNGraphEdge(const drwnNNGraphNodeIndex& tgtIndx, float w) :
        targetNode(tgtIndx), weight(w), status(DRWN_NNG_DIRTY) { /* do nothing */ }

    bool operator<(const drwnNNGraphEdge& e) const {
        if (weight < e.weight) return true;
        if (weight > e.weight) return false;
        return (targetNode < e.targetNode);
    }
};

// drwnNNGraphEdge sorting functions -----------------------------------------

class drwnNNGraphSortByScore : public std::binary_function<drwnNNGraphEdge, drwnNNGraphEdge, bool> {
 public:
    bool operator()(const drwnNNGraphEdge& a, const drwnNNGraphEdge& b) const {
        return (a.weight < b.weight);
    }
};

class drwnNNGraphSortByImage : public std::binary_function<drwnNNGraphEdge, drwnNNGraphEdge, bool> {
 public:
    bool operator()(const drwnNNGraphEdge& a, const drwnNNGraphEdge& b) const {
        return (a.targetNode.imgIndx < b.targetNode.imgIndx) ||
            ((a.targetNode.imgIndx == b.targetNode.imgIndx) && (a.weight < b.weight));
    }
};

// drwnNNGraphEdgeList -------------------------------------------------------

typedef list<drwnNNGraphEdge> drwnNNGraphEdgeList;

// drwnNNGraphNode -----------------------------------------------------------

class drwnNNGraphNode : public drwnPersistentRecord {
 public:
    VectorXf features;                
    int32_t label;                    

    drwnNNGraphEdgeList edges;        
    set<uint16_t> spatialNeighbours;  

 public:
    drwnNNGraphNode() : label(-1) { /* do nothing */ }
    drwnNNGraphNode(const VectorXf& x, int32_t y = -1) :
        features(x), label(-1) { /* do nothing */ }

    void clear();

    bool insert(const drwnNNGraphEdge& e);

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

// drwnNNGraphImageData ------------------------------------------------------

class drwnNNGraphImageData {
 public:
    static string imgDir;              
    static string imgExt;              
    static string lblDir;              
    static string lblExt;              
    static string segDir;              
    static string segExt;              

 protected:
    string _name;                       
    cv::Mat _img;                       
    MatrixXi _labels;                   
    drwnSuperpixelContainer _segments;  
    vector<unsigned> _colours;          
    vector<cv::Point> _centroids;       

 public:
    drwnNNGraphImageData(const string &name);
    drwnNNGraphImageData(const cv::Mat& img, const drwnSuperpixelContainer& segments);

    const string& name() const { return _name; }
    const cv::Mat& image() const { return _img; }
    const MatrixXi& labels() const { return _labels; }
    const drwnSuperpixelContainer& segments() const { return _segments; }
    drwnSuperpixelContainer& segments() { return _segments; }

    size_t numSegments() const { return _segments.size(); }
    size_t height() const { return _img.rows; }
    size_t width() const { return _img.cols; }

    unsigned colour(unsigned segId) const { return _colours[segId]; }
    cv::Scalar rgbColour(unsigned segId) const { return CV_RGB(_colours[segId] & 0x0000ff,
        (_colours[segId] >> 8) & 0x0000ff, (_colours[segId] >> 16) & 0x0000ff); }

    cv::Point centroid(unsigned segId) const { return _centroids[segId]; }

    void setLabels(const MatrixXi& labels);

    vector<VectorXd> getSegmentLabelMarginals(int numLabels = -1) const;

 protected:
    void cacheSegmentData();
};

// drwnNNGraphImage ----------------------------------------------------------

class drwnNNGraphImage : public drwnPersistentRecord {
 public:
    bool bSourceMatchable;

    bool bTargetMatchable;

    int32_t eqvClass;

 protected:
    string _name;                     
    vector<drwnNNGraphNode> _nodes;   

 public:
    drwnNNGraphImage() : drwnPersistentRecord(), bSourceMatchable(true), bTargetMatchable(true),
        eqvClass(-1), _name("") { /* do nothing */ }
    drwnNNGraphImage(const string& name, unsigned n = 0);
    drwnNNGraphImage(const drwnNNGraphImageData& image);
    ~drwnNNGraphImage() { /* do nothing */ }

    const string& name() const { return _name; }
    size_t numNodes() const { return _nodes.size(); }

    void clearNodes();
    void clearEdges();

    void initialize(const string& name, unsigned n = 0);
    virtual void initialize(const drwnNNGraphImageData& image);

    void transformNodeFeatures(const drwnFeatureTransform& xform);
    void appendNodeFeatures(const drwnNNGraphImageData& image, const cv::Mat& features);
    void appendNodeFeatures(const drwnNNGraphImageData& image, const vector<cv::Mat>& features);

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

    drwnNNGraphImage clone(bool bWithFeatures = true) const;

    inline const drwnNNGraphNode& operator[](unsigned segId) const { return _nodes[segId]; }
    inline drwnNNGraphNode& operator[](unsigned segId) { return _nodes[segId]; }

 protected:
    virtual void cacheNodeNeighbourhoods(const drwnNNGraphImageData& image);
    virtual void cacheNodeFeatures(const drwnNNGraphImageData& image);
    virtual void cacheNodeLabels(const drwnNNGraphImageData& image);
};

// drwnNNGraph ---------------------------------------------------------------
class drwnNNGraph {
 public:
    static unsigned int K;             
    static bool DO_PROPAGATE;          
    static bool DO_LOCAL;              
    static bool DO_SEARCH;             
    static int  DO_RANDPROJ;           
    static bool DO_ENRICHMENT;         
    static int DO_EXHAUSTIVE;          

 protected:
    vector<drwnNNGraphImage> _images;  
    map<string, unsigned> _names;      

 public:
    drwnNNGraph() { /* do nothing */ }
    virtual ~drwnNNGraph() { /* do nothing */ }

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

    void clear() { _images.clear(); }
    void reserve(size_t n) { _images.reserve(_images.size() + n); }

    drwnNNGraph clone(bool bWithFeatures = true) const;

    size_t numImages() const { return _images.size(); }
    size_t numNodes() const;
    size_t numEdges() const;

    size_t numNodesWithLabel(int label) const;

    int findImage(const string& baseName) const;
    int appendImage(const string& baseName, unsigned numNodes = 0);
    int appendImage(const drwnNNGraphImage& image);
    int removeImage(unsigned imgIndx);

    pair<double, double> energy() const;

    inline bool inSameEqvClass(unsigned imgIndxA, unsigned imgIndxB) const {
        if (imgIndxA == imgIndxB) return true;
        if (_images[imgIndxA].eqvClass < 0) return false;
        return (_images[imgIndxA].eqvClass == _images[imgIndxB].eqvClass);
    }

    inline const drwnNNGraphImage& operator[](size_t indx) const { return _images[indx]; }
    inline drwnNNGraphImage& operator[](size_t indx) { return _images[indx]; }

    inline const drwnNNGraphImage& operator[](const string& baseName) const {
        const map<string, unsigned>::const_iterator it = _names.find(baseName);
        return _images[it->second]; 
    }
    inline drwnNNGraphImage& operator[](const string& baseName) {
        const map<string, unsigned>::const_iterator it = _names.find(baseName);
        return _images[it->second]; 
    }

    inline const drwnNNGraphNode& operator[](const drwnNNGraphNodeIndex& indx) const {
        return _images[indx.imgIndx][indx.segId];
    }
    inline drwnNNGraphNode& operator[](const drwnNNGraphNodeIndex& indx) {
        return _images[indx.imgIndx][indx.segId];
    }
};

// drwnNNGraphNodeAnnotation -------------------------------------------------
template <typename T>
class drwnNNGraphNodeAnnotation {
 protected:
    vector<vector<T> > _tags;

 public:
    drwnNNGraphNodeAnnotation() { /* do nothing */ }
    drwnNNGraphNodeAnnotation(const drwnNNGraph& graph, const T& initValue) {
        initialize(graph, initValue);
    }
    ~drwnNNGraphNodeAnnotation() { /* do nothing */ }

    void initialize(const drwnNNGraph& graph, const T& initValue) {
        _tags.clear();
        _tags.resize(graph.numImages());
        for (size_t imgIndx = 0; imgIndx < graph.numImages(); imgIndx++) {
            _tags[imgIndx].resize(graph[imgIndx].numNodes(), initValue);
        }
    }

    void reset(const T& initValue) {
        for (size_t imgIndx = 0; imgIndx < _tags.size(); imgIndx++) {
            for (size_t segId = 0; segId < _tags[imgIndx].size(); segId++) {
                _tags[imgIndx][segId] = initValue;
            }
        }
    }

    inline const vector<T>& operator[](unsigned imgIndx) const {
        return _tags[imgIndx];
    }
    inline vector<T>& operator[](unsigned imgIndx) {
        return _tags[imgIndx];
    }

    inline const T& operator[](const drwnNNGraphNodeIndex& indx) const {
        return _tags[indx.imgIndx][indx.segId];
    }
    inline T& operator[](const drwnNNGraphNodeIndex& indx) {
        return _tags[indx.imgIndx][indx.segId];
    }
};