Point Cloud Library (PCL)  1.7.1
grabcut.hpp
1 ///////////////////////////////////////////////////////
2 // grabCut->initialize (xstart, ystart, xend, yend); //
3 // grabCut->fitGMMs (); //
4 ///////////////////////////////////////////////////////
5 #ifndef PCL_SEGMENTATION_IMPL_GRABCUT_HPP
6 #define PCL_SEGMENTATION_IMPL_GRABCUT_HPP
7 
8 #include <pcl/search/organized.h>
9 #include <pcl/search/kdtree.h>
10 #include <pcl/common/distances.h>
11 
12 namespace pcl
13 {
14  template <>
17  {
18  return ((c1.r-c2.r)*(c1.r-c2.r)+(c1.g-c2.g)*(c1.g-c2.g)+(c1.b-c2.b)*(c1.b-c2.b));
19  }
20 }
21 
22 template <typename PointT> void
24 {
25  input_ = cloud;
26 }
27 
28 template <typename PointT> bool
30 {
31  using namespace pcl::segmentation::grabcut;
33  {
34  PCL_ERROR ("[pcl::GrabCut::initCompute ()] Init failed!");
35  return (false);
36  }
37 
38  if (!input_->isOrganized ())
39  {
40  PCL_ERROR ("[pcl::GrabCut::initCompute ()] Need an organized point cloud to proceed!");
41  return (false);
42  }
43 
44  std::vector<pcl::PCLPointField> in_fields_;
45  if ((pcl::getFieldIndex<PointT> (*input_, "rgb", in_fields_) == -1) &&
46  (pcl::getFieldIndex<PointT> (*input_, "rgba", in_fields_) == -1))
47  {
48  PCL_ERROR ("[pcl::GrabCut::initCompute ()] No RGB data available, aborting!");
49  return (false);
50  }
51 
52  // Initialize the working image
53  image_.reset (new Image (input_->width, input_->height));
54  for (std::size_t i = 0; i < input_->size (); ++i)
55  {
56  (*image_) [i] = Color (input_->points[i]);
57  }
58  width_ = image_->width;
59  height_ = image_->height;
60 
61  // Initialize the spatial locator
62  if (!tree_)
63  {
64  if (input_->isOrganized ())
65  tree_.reset (new pcl::search::OrganizedNeighbor<PointT> ());
66  else
67  tree_.reset (new pcl::search::KdTree<PointT> (false));
68  tree_->setInputCloud (input_);
69  }
70  const std::size_t indices_size = indices_->size ();
71  trimap_ = std::vector<segmentation::grabcut::TrimapValue> (indices_size, TrimapUnknown);
72  hard_segmentation_ = std::vector<segmentation::grabcut::SegmentationValue> (indices_size,
74  GMM_component_.resize (indices_size);
75  n_links_.resize (indices_size);
76  // soft_segmentation_ = 0; // Not yet implemented
77  foreground_GMM_.resize (K_);
78  background_GMM_.resize (K_);
79  //set some constants
80  computeL ();
81  computeBeta ();
82  computeNLinks ();
83 
84  initialized_ = false;
85  return (true);
86 }
87 
88 template <typename PointT> void
89 pcl::GrabCut<PointT>::addEdge (vertex_descriptor v1, vertex_descriptor v2, float capacity, float rev_capacity)
90 {
91  graph_.addEdge (v1, v2, capacity, rev_capacity);
92 }
93 
94 template <typename PointT> void
95 pcl::GrabCut<PointT>::setTerminalWeights (vertex_descriptor v, float source_capacity, float sink_capacity)
96 {
97  graph_.addSourceEdge (v, source_capacity);
98  graph_.addTargetEdge (v, sink_capacity);
99 }
100 
101 // template <typename PointT> void
102 // pcl::GrabCut<PointT>::setBackgroundPointsIndices (int x1, int y1, int x2, int y2)
103 // {
104 // using namespace pcl::segmentation::grabcut;
105 
106 // // Step 1: User creates inital Trimap with rectangle, Background outside, Unknown inside
107 // fill (trimap_.begin (), trimap_.end (), TrimapBackground);
108 // fillRectangle (trimap_, width_, height_, x1, y1, x2, y2, TrimapUnknown);
109 
110 // // Step 2: Initial segmentation, Background where Trimap is Background, Foreground where Trimap is Unknown.
111 // fill (hard_segmentation_.begin (), hard_segmentation_.end (), SegmentationBackground);
112 // fillRectangle (hard_segmentation_, width_, height_, x1, y1, x2, y2, SegmentationForeground);
113 // if (!initialized_)
114 // {
115 // fitGMMs ();
116 // initialized_ = true;
117 // }
118 // }
119 
120 template <typename PointT> void
122 {
123  using namespace pcl::segmentation::grabcut;
124  if (!initCompute ())
125  return;
126 
127  std::fill (trimap_.begin (), trimap_.end (), TrimapBackground);
128  std::fill (hard_segmentation_.begin (), hard_segmentation_.end (), SegmentationBackground);
129  for (std::vector<int>::const_iterator idx = indices->indices.begin (); idx != indices->indices.end (); ++idx)
130  {
131  trimap_[*idx] = TrimapUnknown;
132  hard_segmentation_[*idx] = SegmentationForeground;
133  }
134 
135  if (!initialized_)
136  {
137  fitGMMs ();
138  initialized_ = true;
139  }
140 }
141 
142 template <typename PointT> void
144 {
145  // Step 3: Build GMMs using Orchard-Bouman clustering algorithm
146  buildGMMs (*image_, *indices_, hard_segmentation_, GMM_component_, background_GMM_, foreground_GMM_);
147 
148  // Initialize the graph for graphcut (do this here so that the T-Link debugging image will be initialized)
149  initGraph ();
150 }
151 
152 template <typename PointT> int
154 {
155  // Steps 4 and 5: Learn new GMMs from current segmentation
156  learnGMMs (*image_, *indices_, hard_segmentation_, GMM_component_, background_GMM_, foreground_GMM_);
157 
158  // Step 6: Run GraphCut and update segmentation
159  initGraph ();
160 
161  float flow = graph_.solve ();
162 
163  int changed = updateHardSegmentation ();
164  PCL_INFO ("%d pixels changed segmentation (max flow = %f)\n", changed, flow);
165 
166  return (changed);
167 }
168 
169 template <typename PointT> void
171 {
172  std::size_t changed = indices_->size ();
173 
174  while (changed)
175  changed = refineOnce ();
176 }
177 
178 template <typename PointT> int
180 {
181  using namespace pcl::segmentation::grabcut;
182 
183  int changed = 0;
184 
185  const int number_of_indices = static_cast<int> (indices_->size ());
186  for (int i_point = 0; i_point < number_of_indices; ++i_point)
187  {
188  SegmentationValue old_value = hard_segmentation_ [i_point];
189 
190  if (trimap_ [i_point] == TrimapBackground)
191  hard_segmentation_ [i_point] = SegmentationBackground;
192  else
193  if (trimap_ [i_point] == TrimapForeground)
194  hard_segmentation_ [i_point] = SegmentationForeground;
195  else // TrimapUnknown
196  {
197  if (isSource (graph_nodes_[i_point]))
198  hard_segmentation_ [i_point] = SegmentationForeground;
199  else
200  hard_segmentation_ [i_point] = SegmentationBackground;
201  }
202 
203  if (old_value != hard_segmentation_ [i_point])
204  ++changed;
205  }
206  return (changed);
207 }
208 
209 template <typename PointT> void
211 {
212  using namespace pcl::segmentation::grabcut;
213  std::vector<int>::const_iterator idx = indices->indices.begin ();
214  for (; idx != indices->indices.end (); ++idx)
215  trimap_[*idx] = t;
216 
217  // Immediately set the hard segmentation as well so that the display will update.
218  if (t == TrimapForeground)
219  for (idx = indices->indices.begin (); idx != indices->indices.end (); ++idx)
220  hard_segmentation_[*idx] = SegmentationForeground;
221  else
222  if (t == TrimapBackground)
223  for (idx = indices->indices.begin (); idx != indices->indices.end (); ++idx)
224  hard_segmentation_[*idx] = SegmentationBackground;
225 }
226 
227 template <typename PointT> void
229 {
230  using namespace pcl::segmentation::grabcut;
231  const int number_of_indices = static_cast<int> (indices_->size ());
232  // Set up the graph (it can only be used once, so we have to recreate it each time the graph is updated)
233  graph_.clear ();
234  graph_nodes_.clear ();
235  graph_nodes_.resize (indices_->size ());
236  int start = graph_.addNodes (indices_->size ());
237  for (int idx = 0; idx < indices_->size (); ++idx)
238  {
239  graph_nodes_[idx] = start;
240  ++start;
241  }
242 
243  // Set T-Link weights
244  for (int i_point = 0; i_point < number_of_indices; ++i_point)
245  {
246  int point_index = (*indices_) [i_point];
247  float back, fore;
248 
249  switch (trimap_[point_index])
250  {
251  case TrimapUnknown :
252  {
253  fore = static_cast<float> (-log (background_GMM_.probabilityDensity (image_->points[point_index])));
254  back = static_cast<float> (-log (foreground_GMM_.probabilityDensity (image_->points[point_index])));
255  break;
256  }
257  case TrimapBackground :
258  {
259  fore = 0;
260  back = L_;
261  break;
262  }
263  default :
264  {
265  fore = L_;
266  back = 0;
267  }
268  }
269 
270  setTerminalWeights (graph_nodes_[i_point], fore, back);
271  }
272 
273  // Set N-Link weights from precomputed values
274  for (int i_point = 0; i_point < number_of_indices; ++i_point)
275  {
276  const NLinks &n_link = n_links_[i_point];
277  if (n_link.nb_links > 0)
278  {
279  int point_index = (*indices_) [i_point];
280  std::vector<int>::const_iterator indices_it = n_link.indices.begin ();
281  std::vector<float>::const_iterator weights_it = n_link.weights.begin ();
282  for (; indices_it != n_link.indices.end (); ++indices_it, ++weights_it)
283  {
284  if (*indices_it != point_index)
285  {
286  addEdge (graph_nodes_[i_point], graph_nodes_[*indices_it], *weights_it, *weights_it);
287  }
288  }
289  }
290  }
291 }
292 
293 template <typename PointT> void
295 {
296  const int number_of_indices = static_cast<int> (indices_->size ());
297  for (int i_point = 0; i_point < number_of_indices; ++i_point)
298  {
299  NLinks &n_link = n_links_[i_point];
300  if (n_link.nb_links > 0)
301  {
302  int point_index = (*indices_) [i_point];
303  std::vector<int>::const_iterator indices_it = n_link.indices.begin ();
304  std::vector<float>::const_iterator dists_it = n_link.dists.begin ();
305  std::vector<float>::iterator weights_it = n_link.weights.begin ();
306  for (; indices_it != n_link.indices.end (); ++indices_it, ++dists_it, ++weights_it)
307  {
308  if (*indices_it != point_index)
309  {
310  // We saved the color distance previously at the computeBeta stage for optimization purpose
311  float color_distance = *weights_it;
312  // Set the real weight
313  *weights_it = static_cast<float> (lambda_ * exp (-beta_ * color_distance) / sqrt (*dists_it));
314  }
315  }
316  }
317  }
318 }
319 
320 template <typename PointT> void
322 {
323  float result = 0;
324  std::size_t edges = 0;
325 
326  const int number_of_indices = static_cast<int> (indices_->size ());
327 
328  for (int i_point = 0; i_point < number_of_indices; i_point++)
329  {
330  int point_index = (*indices_)[i_point];
331  const PointT& point = input_->points [point_index];
332  if (pcl::isFinite (point))
333  {
334  NLinks &links = n_links_[i_point];
335  int found = tree_->nearestKSearch (point, nb_neighbours_, links.indices, links.dists);
336  if (found > 1)
337  {
338  links.nb_links = found - 1;
339  links.weights.reserve (links.nb_links);
340  for (std::vector<int>::const_iterator nn_it = links.indices.begin (); nn_it != links.indices.end (); ++nn_it)
341  {
342  if (*nn_it != point_index)
343  {
344  float color_distance = squaredEuclideanDistance (image_->points[point_index], image_->points[*nn_it]);
345  links.weights.push_back (color_distance);
346  result+= color_distance;
347  ++edges;
348  }
349  else
350  links.weights.push_back (0.f);
351  }
352  }
353  }
354  }
355  std::cout << "result " << result << std::endl;
356  std::cout << "edges " << edges << std::endl;
357  beta_ = 1e5 / (2*result / edges);
358  std::cout << "beta " << beta_ << std::endl;
359 }
360 
361 template <typename PointT> void
363 {
364  L_ = 8*lambda_ + 1;
365 }
366 
367 template <typename PointT> void
368 pcl::GrabCut<PointT>::extract (std::vector<pcl::PointIndices>& clusters)
369 {
370  using namespace pcl::segmentation::grabcut;
371  clusters.clear ();
372  clusters.resize (2);
373  clusters[0].indices.reserve (indices_->size ());
374  clusters[1].indices.reserve (indices_->size ());
375  refine ();
376  assert (hard_segmentation_.size () == indices_->size ());
377  const int indices_size = static_cast<int> (indices_->size ());
378  for (int i = 0; i < indices_size; ++i)
379  if (hard_segmentation_[i] == SegmentationForeground)
380  clusters[1].indices.push_back (i);
381  else
382  clusters[0].indices.push_back (i);
383 }
384 
385 #endif