drwnGrabCutInstance.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:    drwnGrabCutInstance.h
** AUTHOR(S):   Stephen Gould <stephen.gould@anu.edu.au>
**      Kevin Guo <Kevin.Guo@nicta.com.au>
**
*****************************************************************************/

#pragma once
#include "Eigen/Core"

#include "cv.h"

#include "drwnBase.h"
#include "drwnIO.h"
#include "drwnML.h"

#include "drwnColourHistogram.h"
#include "drwnPixelNeighbourContrasts.h"

using namespace std;
using namespace Eigen;

// drwnGrabCutInstance ------------------------------------------------------

class drwnGrabCutInstance {
public:
    static const unsigned char MASK_FG = 0xff;       
    static const unsigned char MASK_BG = 0x00;       
    static const unsigned char MASK_C_FG = 0x80;     
    static const unsigned char MASK_C_BG = 0x40;     
    static const unsigned char MASK_C_BOTH = 0xc0;   
    static const unsigned char MASK_C_NONE = 0x20;   

    static bool bVisualize;    
    static int maxIterations;  

public:
    // meta parameters
    string name;               

protected:
    // instance definition
    cv::Mat _img;              
    cv::Mat _trueMask;         
    cv::Mat _mask;             
    int _numUnknown;           

    // cached data
    cv::Mat _unary;            
    drwnPixelNeighbourContrasts *_pairwise; 

    // model parameters
    double _unaryWeight;       
    double _pottsWeight;       
    double _pairwiseWeight;    

public:
    drwnGrabCutInstance();
    drwnGrabCutInstance(const drwnGrabCutInstance& instance);
    virtual ~drwnGrabCutInstance();

    int width() const { return _img.cols; }
    int height() const { return _img.rows; }
    int size() const { return _img.cols * _img.rows; }
    int numUnknown() const { return _numUnknown; }

    // caller does not need to free returned objects
    const cv::Mat& image() const { return _img; }
    const cv::Mat& trueSegmentation() const { return _trueMask; }
    const cv::Mat& segmentationMask() const { return _mask; }

    cv::Mat knownForeground() const;
    cv::Mat knownBackground() const;
    cv::Mat unknownPixels() const;
    cv::Mat foregroundColourMask() const;
    cv::Mat backgroundColourMask() const;

    bool isUnknownPixel(int x, int y, const cv::Mat& mask) const {
        const unsigned char p = mask.at<unsigned char>(y, x);
        return ((p != MASK_FG) && (p != MASK_BG));
    }
    bool isUnknownPixel(int x, int y) const {
        return isUnknownPixel(x, y, _mask);
    }

    void initialize(const cv::Mat& img, const cv::Rect& rect, const char *colorModelFile = NULL);
    void initialize(const cv::Mat& img, const cv::Mat& inferMask, const char *colorModelFile = NULL);
    void initialize(const cv::Mat& img, const cv::Rect& rect, const cv::Mat& trueMask,
        const char *colorModelFile = NULL);
    void initialize(const cv::Mat& img, const cv::Mat& inferMask, const cv::Mat& trueMask,
        const char *colorModelFile = NULL);

    virtual void loadColourModels(const char *filename) = 0;
    virtual void saveColourModels(const char *filename) const = 0;

    void setBaseModelWeights(double u, double p, double c);

    const cv::Mat& unaryPotentials() const { return _unary; }

    const drwnPixelNeighbourContrasts *pairwisePotentials() { return _pairwise; }

    double unaryEnergy(const cv::Mat& seg) const;

    double pottsEnergy(const cv::Mat& seg) const;

    double pairwiseEnergy(const cv::Mat& seg) const;

    virtual double energy(const cv::Mat& seg) const;

    double foregroundRatio(const cv::Mat& seg) const;

    inline double backgroundRatio(const cv::Mat& seg) const {
        return 1.0 - foregroundRatio(seg);
    }

    double loss(const cv::Mat& seg) const;

    virtual cv::Mat inference();
    virtual cv::Mat lossAugmentedInference();

    virtual cv::Mat visualize(const cv::Mat& seg) const;

protected:
    void free();

    inline vector<double> pixelColour(int y, int x) const;

    virtual void learnColourModel(const cv::Mat& mask, bool bForeground) = 0;

    virtual void updateUnaryPotentials() = 0;

    virtual cv::Mat graphCut(const cv::Mat& unary) const;
};

// drwnGrabCutInstanceGMM ---------------------------------------------------

class drwnGrabCutInstanceGMM : public drwnGrabCutInstance {
 public:
    static size_t maxSamples;  
    static int numMixtures;    

 protected:
    drwnGaussianMixture _fgColourModel; 
    drwnGaussianMixture _bgColourModel; 

 public:
    drwnGrabCutInstanceGMM();
    virtual ~drwnGrabCutInstanceGMM();

    void loadColourModels(const char *filename);
    void saveColourModels(const char *filename) const;

protected:
    void learnColourModel(const cv::Mat& mask, bool bForeground);

    void updateUnaryPotentials();
};

// drwnGrabCutInstanceHistogram ---------------------------------------------

class drwnGrabCutInstanceHistogram : public drwnGrabCutInstance {
 public:
    static double pseudoCounts;   
    static unsigned channelBits;  

 protected:
    drwnColourHistogram _fgColourModel; 
    drwnColourHistogram _bgColourModel; 

 public:
    drwnGrabCutInstanceHistogram();
    virtual ~drwnGrabCutInstanceHistogram();

    void loadColourModels(const char *filename);
    void saveColourModels(const char *filename) const;

protected:
    void learnColourModel(const cv::Mat& mask, bool bForeground);

    void updateUnaryPotentials();
};