image_segmentation_demo: image_segmentation_demo: grabcut3d_object

00001 /*********************************************************************
00002  *
00003  * Software License Agreement (BSD License)
00004  *
00005  *  Copyright (c) 2011, Robert Bosch LLC.
00006  *  All rights reserved.
00007  *
00008  *  Redistribution and use in source and binary forms, with or without
00009  *  modification, are permitted provided that the following conditions
00010  *  are met:
00011  *
00012  *   * Redistributions of source code must retain the above copyright
00013  *     notice, this list of conditions and the following disclaimer.
00014  *   * Redistributions in binary form must reproduce the above
00015  *     copyright notice, this list of conditions and the following
00016  *     disclaimer in the documentation and/or other materials provided
00017  *     with the distribution.
00018  *   * Neither the name of the Robert Bosch nor the names of its
00019  *     contributors may be used to endorse or promote products derived
00020  *     from this software without specific prior written permission.
00021  *
00022  *  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
00023  *  "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
00024  *  LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
00025  *  FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
00026  *  COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
00027  *  INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
00028  *  BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
00029  *  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
00030  *  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
00031  *  LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
00032  *  ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
00033  *  POSSIBILITY OF SUCH DAMAGE.
00034  *
00035  *********************************************************************/
00036 
00037 #include "image_segmentation_demo/grabcut3d_object_segmenter.h"
00038 #include "grabcut_3d/grabcut_3d.h"
00039 
00040 using namespace std;
00041 using namespace cv;
00042 
00043 namespace image_segmentation_demo{
00044 
00045 static const Scalar RD = Scalar(0,0,255);
00046 static const Scalar PK = Scalar(230,130,255);
00047 static const Scalar BL = Scalar(255,0,0);
00048 static const Scalar LB = Scalar(255,255,160);
00049 static const Scalar GR = Scalar(0,255,0);
00050 static const Scalar WH = Scalar(255,255,255);
00051 static const Scalar GY = Scalar(128,128,128);
00052 static const Scalar BK = Scalar(0,0,0);
00053 static const int DOT_RADIUS = 2;
00054 static const int LINE_THICKNESS = -1;
00055 
00056 // Initialize default rectangle marker (not actually used as a rectangle)
00057 const Rect GrabCut3DObjectSegmenter::DEFAULT_RECT(0,0,0,0);
00058 // This scale factor is what is actually used to calculate the default
00059 // rectangle.
00060 const int DEFAULT_RECT_SCALE_FACTOR = 2;
00061 
00062 /* Helper function to convert mask in 4-value GrabCut format to a binary mask
00063  * with foreground = 1, background = 0. */
00064 static void
00065 binaryMaskFromGCMask(const Mat &_gc_mask, Mat &_bin_mask)
00066 {
00067     if( _gc_mask.empty() || _gc_mask.type() != CV_8UC1 )
00068         CV_Error( CV_StsBadArg, "_gc_mask is empty or has incorrect type (not CV_8UC1)" );
00069 
00070     if(    _bin_mask.empty()
00071         || _bin_mask.rows!=_gc_mask.rows
00072         || _bin_mask.cols!=_gc_mask.cols )
00073         _bin_mask.create( _gc_mask.size(), CV_8UC1 );
00074 
00075     // Convert GC_BGD and GC_PR_BGD to 0, GC_FGD and GC_PR_FGD to 1
00076     _bin_mask = _gc_mask & 1;
00077 }
00078 
00079 GrabCut3DObjectSegmenter::GrabCut3DObjectSegmenter()
00080 {
00081   // Initialize window
00082   //setImages(_img,_depth_image);
00083   updateDisplayImage();
00084 
00085   iter_count_ = 0;
00086 }
00087 
00088 GrabCut3DObjectSegmenter::~GrabCut3DObjectSegmenter()
00089 {
00090 
00091 }
00092 
00093 /* Mutator for initialized state of GrabCut3DObjectSegmenter object. Client can set
00094  * initialized=true with no side effects; setting initialized=false clears
00095  * background and foreground pixel arrays and returns application state
00096  * variables to NOT_SET. */
00097 // TODO: Should this be public?
00098 void
00099 GrabCut3DObjectSegmenter::initializedIs(bool _init)
00100 {
00101   if (_init)
00102     initialized_=_init;
00103   else
00104   {
00105     if ( !mask_.empty() ) mask_.setTo(Scalar::all(GC_BGD));
00106     bgd_pxls_.clear();   fgd_pxls_.clear();
00107     pr_bgd_pxls_.clear();  pr_fgd_pxls_.clear();
00108     initialized_   = _init;
00109     rect_state_ = NOT_SET;
00110     lbls_state_ = NOT_SET;
00111     pr_lbls_state_ = NOT_SET;
00112     iterCountIs(0);
00113     updateDisplayImage();
00114   }
00115 }
00116 
00117 /* Mutator for stored image. */
00118 void GrabCut3DObjectSegmenter::setImages(const Mat& _image, const Mat& _depth_image)
00119 {
00120     if( _image.empty() ) {
00121       ROS_WARN("GrabCut3DObjectSegmenter: setImages called with empty image");
00122       return;
00123     }
00124     if (_depth_image.empty()){
00125         ROS_WARN("GrabCut3DObjectSegmenter: setImages called with empty depth image");
00126         return;
00127     }
00128     ROS_INFO_STREAM("image size " << _image.size().width << " "<<_image.size().height);
00129     image_ = _image.clone();
00130     ROS_INFO_STREAM("depth image size " << _depth_image.size().width << " "<<_depth_image.size().height);
00131     depth_image_ = _depth_image.clone();
00132     display_image_ = image_.clone();
00133 
00134     mask_.create( image_.size(), CV_8UC1);
00135     initializedIs(false);
00136 }
00137 
00138 /* Mutator for rectangle. */
00139 void GrabCut3DObjectSegmenter::rectIs(const Rect &_r)
00140 {
00141   if (_r == DEFAULT_RECT)
00142     rect_ = Rect(image_.cols/(DEFAULT_RECT_SCALE_FACTOR*2),
00143                  image_.rows/(DEFAULT_RECT_SCALE_FACTOR*2),
00144                  image_.cols/DEFAULT_RECT_SCALE_FACTOR,
00145                  image_.rows/DEFAULT_RECT_SCALE_FACTOR);
00146   else
00147     rect_=_r;
00148 
00149   rect_state_ = SET;
00150 
00151   ROS_INFO_STREAM("rect area is " << rect_.area());
00152   if (rect_.area()==0)
00153         rect_state_=EMPTY;
00154 
00155   setRectInMask();
00156   assert( bgd_pxls_.empty()
00157   && fgd_pxls_.empty()
00158   && pr_bgd_pxls_.empty()
00159   && pr_fgd_pxls_.empty() );
00160   updateDisplayImage();
00161 
00162 }
00163 
00164 /* Accessor for binary mask. */
00165 Mat GrabCut3DObjectSegmenter::binaryMask()
00166 {
00167   Mat bin_mask;
00168   /*
00169   cv::imwrite("/home/stm1pal/ros/stacks/bosch-ros-pkg/mas/grabcut_app/full_mask.pbm", mask_);
00170   double fm_max;
00171   double fm_min;
00172   minMaxLoc(mask_, &fm_min, &fm_max);
00173   std::cout << "fm_max: " << fm_max << std::endl;
00174   std::cout << "fm_min: " << fm_min << std::endl;
00175   */
00176 
00177   binaryMaskFromGCMask(mask_, bin_mask);
00178   /*
00179   cv::imwrite("/home/stm1pal/ros/stacks/bosch-ros-pkg/mas/grabcut_app/bin_mask.pbm", bin_mask);
00180   double bm_max;
00181   double bm_min;
00182   minMaxLoc(bin_mask, &bm_min, &bm_max);
00183   std::cout << "bm_max: " << bm_max << std::endl;
00184   std::cout << "bm_min: " << bm_min << std::endl;
00185   */
00186   return bin_mask;
00187 }
00188 
00189 /* Mutator for window background color. */
00190 void GrabCut3DObjectSegmenter::setWinColor(WinColor _c)
00191 {
00192   win_color_=_c;
00193   updateDisplayImage();
00194 }
00195 
00196 void GrabCut3DObjectSegmenter::updateDisplayImage()
00197 {
00198 
00199   if( image_.empty() )
00200   {
00201     ROS_WARN("GrabCut3DObjectSegmenter: window image state updated with empty image");
00202     return;
00203   }
00204 
00205   // Generate the display image
00206   Scalar color = BK;
00207   if (win_color_ == WHITE) color = WH;
00208   else if (win_color_ == GRAY) color = GY;
00209   else if (win_color_ == GREEN) color = GR;
00210   else if (win_color_ == BLUE) color = BL;
00211 
00212   display_image_.setTo(color);
00213   Mat bin_mask;
00214   if ( !initialized_ )
00215     // If we haven't created a mask, just copy the base image
00216     image_.copyTo( display_image_ );
00217   else
00218   {
00219     binaryMaskFromGCMask( mask_, bin_mask );
00220     image_.copyTo( display_image_, bin_mask );
00221   }
00222 
00223   // Overlay the marked points, BG and FG, on the display image
00224   vector<Point>::const_iterator it;
00225   for( it = bgd_pxls_.begin(); it != bgd_pxls_.end(); ++it )
00226     circle( display_image_, *it, DOT_RADIUS, BL, LINE_THICKNESS );
00227   for( it = fgd_pxls_.begin(); it != fgd_pxls_.end(); ++it )
00228     circle( display_image_, *it, DOT_RADIUS, RD, LINE_THICKNESS );
00229   for( it = pr_bgd_pxls_.begin(); it != pr_bgd_pxls_.end(); ++it )
00230     circle( display_image_, *it, DOT_RADIUS, LB, LINE_THICKNESS );
00231   for( it = pr_fgd_pxls_.begin(); it != pr_fgd_pxls_.end(); ++it )
00232     circle( display_image_, *it, DOT_RADIUS, PK, LINE_THICKNESS );
00233 
00234   // Add the rectangle
00235   if( rect_state_ == IN_PROCESS || rect_state_ == SET || rect_state_==EMPTY)
00236     rectangle( display_image_,
00237         Point( rect_.x, rect_.y ),
00238         Point(rect_.x + rect_.width, rect_.y + rect_.height ),
00239         GR, 2);
00240 }
00241 
00242 /* Private helper method to transfer pixels from designated rectangle to
00243  * Grabcut mask. */
00244 void GrabCut3DObjectSegmenter::setRectInMask()
00245 {
00246     assert( !mask_.empty() );
00247     mask_.setTo( GC_BGD );
00248     rect_.x = max(0, rect_.x);
00249     rect_.y = max(0, rect_.y);
00250     rect_.width = min(rect_.width, image_.cols-rect_.x);
00251     rect_.height = min(rect_.height, image_.rows-rect_.y);
00252     (mask_(rect_)).setTo( Scalar(GC_PR_FGD) );
00253 }
00254 
00255 /* Private helper method to transfer labels designated in UI to mask. */
00256 void GrabCut3DObjectSegmenter::setLblsInMask(bool control_down, bool shift_down, Point p, bool isPr )
00257 {
00258     vector<Point> *bpxls, *fpxls;
00259     uchar bvalue, fvalue;
00260     if( !isPr )
00261     {
00262       bpxls = &bgd_pxls_;
00263       fpxls = &fgd_pxls_;
00264       bvalue = GC_BGD;
00265       fvalue = GC_FGD;
00266     }
00267     else
00268     {
00269       bpxls = &pr_bgd_pxls_;
00270       fpxls = &pr_fgd_pxls_;
00271       bvalue = GC_PR_BGD;
00272       fvalue = GC_PR_FGD;
00273     }
00274 
00275     if( control_down )
00276     {
00277       bpxls->push_back(p);
00278       circle( mask_, p, DOT_RADIUS, bvalue, LINE_THICKNESS );
00279     }
00280     if( shift_down )
00281     {
00282       fpxls->push_back(p);
00283       circle( mask_, p, DOT_RADIUS, fvalue, LINE_THICKNESS );
00284     }
00285 }
00286 
00287 /* Mouse callback function passed to OpenCV window. Handles marking the ROI
00288  * rectangle and foreground and background pixel hints. */
00289 void GrabCut3DObjectSegmenter::mouseClick( MouseEvent event, int x, int y, bool control_down, bool shift_down)
00290 {
00291     // TODO add bad args check
00292   switch( event )
00293   {
00294     case LEFT_BUTTON_DOWN: // set rect or GC_BGD(GC_FGD) labels
00295       {
00296         bool isb = control_down;
00297         bool isf = shift_down;
00298         if( rect_state_ == NOT_SET && !isb && !isf )
00299         {
00300           rect_state_ = IN_PROCESS;
00301           rect_ = Rect( x, y, 1, 1 );
00302         }
00303         if ( (isb || isf) && rect_state_ == SET )
00304           lbls_state_ = IN_PROCESS;
00305       }
00306       break;
00307 
00308     case RIGHT_BUTTON_DOWN: // set GC_PR_BGD(GC_PR_FGD) labels
00309       {
00310         bool isb = control_down;
00311         bool isf = shift_down;
00312         if ( (isb || isf) && rect_state_ == SET )
00313           pr_lbls_state_ = IN_PROCESS;
00314       }
00315       break;
00316 
00317     case LEFT_BUTTON_UP:
00318       if( rect_state_ == IN_PROCESS )
00319       {
00320         rect_ = Rect( Point(rect_.x, rect_.y), Point(x,y) );
00321         rect_state_ = SET;
00322 
00323         ROS_INFO_STREAM("rect area is " << rect_.area());
00324         if (rect_.area()==0)
00325          rect_state_=EMPTY;
00326 
00327         setRectInMask();
00328         assert( bgd_pxls_.empty()
00329                && fgd_pxls_.empty()
00330                && pr_bgd_pxls_.empty()
00331                && pr_fgd_pxls_.empty() );
00332         updateDisplayImage();
00333       }
00334       if( lbls_state_ == IN_PROCESS )
00335       {
00336         setLblsInMask(control_down, shift_down, Point(x,y), false);
00337         lbls_state_ = SET;
00338         updateDisplayImage();
00339       }
00340       break;
00341 
00342     case RIGHT_BUTTON_UP:
00343       if( pr_lbls_state_ == IN_PROCESS )
00344       {
00345         setLblsInMask(control_down, shift_down, Point(x,y), true);
00346         pr_lbls_state_ = SET;
00347         updateDisplayImage();
00348       }
00349       break;
00350 
00351     case MOUSE_MOVE:
00352       if( rect_state_ == IN_PROCESS )
00353       {
00354         rect_ = Rect( Point(rect_.x, rect_.y), Point(x,y) );
00355         assert( bgd_pxls_.empty()
00356                && fgd_pxls_.empty()
00357                && pr_bgd_pxls_.empty()
00358                && pr_fgd_pxls_.empty() );
00359         updateDisplayImage();
00360       }
00361       else if( lbls_state_ == IN_PROCESS )
00362       {
00363         setLblsInMask(control_down, shift_down, Point(x,y), false);
00364         updateDisplayImage();
00365       }
00366       else if( pr_lbls_state_ == IN_PROCESS )
00367       {
00368         setLblsInMask(control_down, shift_down, Point(x,y), true);
00369         updateDisplayImage();
00370       }
00371       break;
00372   }
00373 }
00374 
00375 void GrabCut3DObjectSegmenter::iterCountIs(int _icnt)
00376 {
00377   if (_icnt == 0) {
00378     iter_count_ = 0;
00379     return;
00380   }
00381 
00382   int iter_inc = _icnt - iter_count_;
00383   if (iter_inc <= 0) return;
00384 
00385   if( initialized() )
00386     grabCut3D( image_, depth_image_, mask_, rect_, bgd_model_, fgd_model_, iter_inc );
00387   else
00388   {
00389     if( rect_state_ != SET ) return;
00390         
00391     if (rect_.area()==0) return;
00392         //  rect_state_=EMPTY;
00393 
00394     if( lbls_state_ == SET || pr_lbls_state_ == SET )
00395       grabCut3D( image_, depth_image_, mask_, rect_, bgd_model_, fgd_model_,
00396                iter_inc, GC_INIT_WITH_MASK );
00397     else
00398       grabCut3D( image_, depth_image_, mask_, rect_, bgd_model_, fgd_model_,
00399                iter_inc, GC_INIT_WITH_RECT );
00400 
00401     initializedIs(true);
00402   }
00403   iter_count_ = _icnt;
00404 
00405   bgd_pxls_.clear(); fgd_pxls_.clear();
00406   pr_bgd_pxls_.clear(); pr_fgd_pxls_.clear();
00407 
00408   updateDisplayImage();
00409 
00410   return;
00411 }
00412 
00413 }
00414
grabcut3d_object_segmenter.cpp
