grabcut.h

Go to the documentation of this file.

/*
 * Software License Agreement (BSD License)
 *
 *  Point Cloud Library (PCL) - www.pointclouds.org
 *  Copyright (c) 2012-, Open Perception, Inc.
 *
 *  All rights reserved.
 *
 *  Redistribution and use in source and binary forms, with or without
 *  modification, are permitted provided that the following conditions
 *  are met:
 *
 *   * Redistributions of source code must retain the above copyright
 *     notice, this list of conditions and the following disclaimer.
 *   * Redistributions in binary form must reproduce the above
 *     copyright notice, this list of conditions and the following
 *     disclaimer in the documentation and/or other materials provided
 *     with the distribution.
 *   * Neither the name of Willow Garage, Inc. nor the names of its
 *     contributors may be used to endorse or promote products derived
 *     from this software without specific prior written permission.
 *
 *  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 *  "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 *  LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 *  FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
 *  COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 *  INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 *  BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 *  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 *  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 *  LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
 *  ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 *  POSSIBILITY OF SUCH DAMAGE.
 *
 * $Id$
 *
 */

#ifndef PCL_SEGMENTATION_GRABCUT
#define PCL_SEGMENTATION_GRABCUT

#include <pcl/point_cloud.h>
#include <pcl/pcl_base.h>
#include <pcl/point_types.h>
#include <pcl/segmentation/boost.h>
#include <pcl/search/search.h>

namespace pcl
{
  namespace segmentation
  {
    namespace grabcut
    {
      class BoykovKolmogorov
      {        
        public:
          typedef int vertex_descriptor;
          typedef double edge_capacity_type;
        
          BoykovKolmogorov (std::size_t max_nodes = 0);
          virtual ~BoykovKolmogorov () {}
          size_t 
          numNodes () const { return nodes_.size (); }
          void 
          reset ();
          void 
          clear ();
          int 
          addNodes (std::size_t n = 1);
          void 
          addConstant (double c) { flow_value_ += c; }
          void 
          addSourceEdge (int u, double cap);
          void 
          addTargetEdge (int u, double cap);
          void 
          addEdge (int u, int v, double cap_uv, double cap_vu = 0.0);
          double
          solve ();
          bool 
          inSourceTree (int u) const { return (cut_[u] == SOURCE); }
          bool 
          inSinkTree (int u) const { return (cut_[u] == TARGET); }
          double
          operator() (int u, int v) const;

        protected:
          typedef enum { FREE = 0x00, SOURCE = 0x01, TARGET = 0x02 } nodestate;
          typedef std::map<int, double> capacitated_edge;
          typedef std::pair<capacitated_edge::iterator, capacitated_edge::iterator> edge_pair;
          void 
          preAugmentPaths ();
          void 
          initializeTrees ();
          std::pair<int, int> 
          expandTrees ();
          void 
          augmentPath (const std::pair<int, int>& path, std::deque<int>& orphans);
          void 
          adoptOrphans (std::deque<int>& orphans);
          void clearActive ();
          inline bool 
          isActiveSetEmpty () const { return (active_head_ == TERMINAL); }
          inline bool 
          isActive (int u) const { return ((u == active_head_) || (active_list_[u].first != TERMINAL)); }
          void 
          markActive (int u);
          void 
          markInactive (int u);
          std::vector<double> source_edges_; 
          std::vector<double> target_edges_; 
          std::vector<capacitated_edge> nodes_; 
          double flow_value_;
          std::vector<unsigned char> cut_;  

        private:
          static const int TERMINAL = -1;
          std::vector<std::pair<int, edge_pair> > parents_; 
          std::vector<std::pair<int, int> > active_list_; 
          int active_head_, active_tail_;
      };

      struct Color
      {
        Color () : r (0), g (0), b (0) {}
        Color (float _r, float _g, float _b) : r(_r), g(_g), b(_b) {}
        Color (const pcl::RGB& color) : r (color.r), g (color.g), b (color.b) {}
        
        template<typename PointT>
        Color (const PointT& p)
        {
          r = static_cast<float> (p.r);
          g = static_cast<float> (p.g);
          b = static_cast<float> (p.b);
        }
        
        template<typename PointT> 
        operator PointT () const
        { 
          PointT p;
          p.r = static_cast<uint32_t> (r);
          p.g = static_cast<uint32_t> (g);
          p.b = static_cast<uint32_t> (b);
          return (p);
        }
        
        float r, g, b;
      };
      typedef pcl::PointCloud<Color> Image;
      float 
      colorDistance (const Color& c1, const Color& c2);
      enum TrimapValue { TrimapUnknown = -1, TrimapForeground, TrimapBackground };
      enum SegmentationValue { SegmentationForeground = 0, SegmentationBackground };      
      struct Gaussian
      {
        Gaussian () {}
        Color mu;
        Eigen::Matrix3f covariance;
        float determinant;
        Eigen::Matrix3f inverse;                        
        float pi;
        float eigenvalue;
        Eigen::Vector3f eigenvector;
      };
      
      class GMM
      {
        public:
          GMM () : gaussians_ (0) {}
          GMM (std::size_t K) : gaussians_ (K) {}
          ~GMM () {}
          std::size_t 
          getK () const { return gaussians_.size (); }
          void 
          resize (std::size_t K) { gaussians_.resize (K); }
          Gaussian& 
          operator[] (std::size_t pos) { return (gaussians_[pos]); }
          const Gaussian& 
          operator[] (std::size_t pos) const { return (gaussians_[pos]); }
          float 
          probabilityDensity (const Color &c);
          float 
          probabilityDensity(std::size_t i, const Color &c);
      
        private:
          std::vector<Gaussian> gaussians_;      
      };

      class GaussianFitter
      {
        public:
        GaussianFitter (float epsilon = 0.0001)
          : sum_ (Eigen::Vector3f::Zero ())
          , accumulator_ (Eigen::Matrix3f::Zero ())
          , count_ (0)
          , epsilon_ (epsilon)
        { }
      
        void 
        add (const Color &c);
        void 
        fit (Gaussian& g, std::size_t total_count, bool compute_eigens = false) const;
        float
        getEpsilon () { return (epsilon_); }
        void
        setEpsilon (float epsilon) { epsilon_ = epsilon; }
      
        private:
        Eigen::Vector3f sum_;
        Eigen::Matrix3f accumulator_;
        uint32_t count_;        
        float epsilon_;
      };

      void 
      buildGMMs (const Image &image, 
                 const std::vector<int>& indices,
                 const std::vector<SegmentationValue> &hardSegmentation,
                 std::vector<std::size_t> &components,
                 GMM &background_GMM, GMM &foreground_GMM);
      void 
      learnGMMs (const Image& image, 
                 const std::vector<int>& indices,
                 const std::vector<SegmentationValue>& hard_segmentation,
                 std::vector<std::size_t>& components,
                 GMM& background_GMM, GMM& foreground_GMM);
    }
  };
  
  template <typename PointT>
  class GrabCut : public pcl::PCLBase<PointT>
  {
    public:
      typedef typename pcl::search::Search<PointT> KdTree;
      typedef typename pcl::search::Search<PointT>::Ptr KdTreePtr;
      typedef typename PCLBase<PointT>::PointCloudConstPtr PointCloudConstPtr;
      typedef typename PCLBase<PointT>::PointCloudPtr PointCloudPtr;
      using PCLBase<PointT>::input_;
      using PCLBase<PointT>::indices_;
      using PCLBase<PointT>::fake_indices_;
    
      GrabCut (uint32_t K = 5, float lambda = 50.f)
        : K_ (K)
        , lambda_ (lambda)
        , initialized_ (false)
        , nb_neighbours_ (9)
      {}
      virtual ~GrabCut () {};
      // /// Set input cloud
      void
      setInputCloud (const PointCloudConstPtr& cloud);
      void
      setBackgroundPoints (const PointCloudConstPtr& background_points);
      void
      setBackgroundPointsIndices (int x1, int y1, int x2, int y2);
      void
      setBackgroundPointsIndices (const PointIndicesConstPtr& indices);
      virtual void 
      refine ();
      virtual int 
      refineOnce ();
      float
      getLambda () { return (lambda_); }
      void
      setLambda (float lambda) { lambda_ = lambda; }
      uint32_t
      getK () { return (K_); }
      void
      setK (uint32_t K) { K_ = K; }
      inline void
      setSearchMethod (const KdTreePtr &tree) { tree_ = tree; }
      inline KdTreePtr
      getSearchMethod () { return (tree_); }
      void
      setNumberOfNeighbours (int nb_neighbours) { nb_neighbours_ = nb_neighbours; }
      int
      getNumberOfNeighbours () const { return (nb_neighbours_); }
      void
      extract (std::vector<pcl::PointIndices>& clusters);

    protected:
      // Storage for N-link weights, each pixel stores links to nb_neighbours      
      struct NLinks
      {
        NLinks () : nb_links (0), indices (0), dists (0), weights (0) {}
        
        int nb_links;
        std::vector<int> indices;
        std::vector<float> dists;
        std::vector<float> weights;
      };
      bool
      initCompute ();    
      typedef pcl::segmentation::grabcut::BoykovKolmogorov::vertex_descriptor vertex_descriptor;
      // /** Update hard segmentation after running GraphCut, \return the number of pixels that have
      //   * changed from foreground to background or vice versa. 
      //   */
      // int 
      // updateHardSegmentation ();
      void 
      computeBeta ();
      void 
      computeL ();
      void 
      computeNLinks ();
      float 
      computeNLink (uint32_t x1, uint32_t y1, uint32_t x2, uint32_t y2);
      void 
      setTrimap (const PointIndicesConstPtr &indices, segmentation::grabcut::TrimapValue t);
      int
      updateHardSegmentation ();
      virtual void 
      fitGMMs ();
      void 
      initGraph ();
      void
      addEdge (vertex_descriptor v1, vertex_descriptor v2, float capacity, float rev_capacity);
      void
      setTerminalWeights (vertex_descriptor v, float source_capacity, float sink_capacity);
      inline bool
      isSource (vertex_descriptor v) { return (graph_.inSourceTree (v)); }
      uint32_t width_;
      uint32_t height_;
      // Variables used in formulas from the paper.
      uint32_t K_;
      float lambda_;            
      float beta_;
      float L_;
      KdTreePtr tree_;
      int nb_neighbours_;
      bool initialized_;
      std::vector<NLinks> n_links_;
      segmentation::grabcut::Image::Ptr image_;
      std::vector<segmentation::grabcut::TrimapValue> trimap_;
      std::vector<std::size_t> GMM_component_;
      std::vector<segmentation::grabcut::SegmentationValue> hard_segmentation_;
      // Not yet implemented (this would be interpreted as alpha)
      std::vector<float> soft_segmentation_;    
      segmentation::grabcut::GMM background_GMM_, foreground_GMM_;
      // Graph part
      pcl::segmentation::grabcut::BoykovKolmogorov graph_;
      std::vector<vertex_descriptor> graph_nodes_;
  };
}

#include <pcl/segmentation/impl/grabcut.hpp>

#endif