Extending Justin Talbot's GrabCut Impl [w/ code]

Justin Talbot has done a tremendous job implementing the GrabCut algorithm in C [link to paper, link to code]. I was missing though, the option to load ANY kind of file, not just PPMs and PGMs.
So I tweaked the code a bit to receive a filename and determine how to load it: use the internal P[P|G]M loaders, or offload the work to the OpenCV image loaders that take in many more type. If the OpenCV method is used, the IplImage is converted to the internal GrabCut code representation.

Image<Color>* load( std::string file_name )
{
 if( file_name.find( ".pgm" ) != std::string::npos )
 {
 return loadFromPGM( file_name );
 }
 else if( file_name.find( ".ppm" ) != std::string::npos )
 {
 return loadFromPPM( file_name );
 }
 else
 {
 return loadOpenCV(file_name);
 }
}
void fromImageMaskToIplImage(const Image<Real>* image, IplImage* ipli) {
 for(int x=0;x<image->width();x++) {
 for(int y=0;y<image->height();y++) {
 //Color c = (*image)(x,y);
 Real r = (*image)(x,y);
 CvScalar s = cvScalarAll(0);
 if(r == 0.0) {
 s.val[0] = 255.0;
 }
 cvSet2D(ipli,ipli->height - y - 1,x,s);
 }
 }
}
Image<Color>* loadIplImage(IplImage* im) {
 Image<Color>* image = new Image<Color>(im->width, im->height);
 for(int x=0;x<im->width;x++) {
 for(int y=0;y<im->height;y++) {
 CvScalar v = cvGet2D(im,im->height-y-1,x);
 Real R, G, B;
 R = (Real)((unsigned char)v.val[2])/255.0f;
 G = (Real)((unsigned char)v.val[1])/255.0f;
 B = (Real)((unsigned char)v.val[0])/255.0f;
 (*image)(x,y) = Color(R,G,B);
 }
 }
 return image;
}
Image<Color>* loadOpenCV(std::string file_name) {
 IplImage* im = cvLoadImage(file_name.c_str(),1);
 Image<Color>* i = loadIplImage(im);
 cvReleaseImage(&im);
 return i;
}

Well, there’s nothing fancy here, but it does give you a fully working GrabCut implementation on top of OpenCV… so there’s the contribution.

GrabCutNS::Image<GrabCutNS::Color>* imageGC = GrabCutNS::loadIplImage(orig);
 GrabCutNS::Image<GrabCutNS::Color>* maskGC = GrabCutNS::loadIplImage(mask);
 GrabCutNS::GrabCut *grabCut = new GrabCutNS::GrabCut( imageGC );
 grabCut->initializeWithMask(maskGC);
 grabCut->fitGMMs();
 //grabCut->refineOnce();
 grabCut->refine();
 IplImage* __GCtmp = cvCreateImage(cvSize(orig->width,orig->height),8,1);
 GrabCutNS::fromImageMaskToIplImage(grabCut->getAlphaImage(),__GCtmp);
 //cvShowImage("result",image);
 cvShowImage("tmp",__GCtmp);
 cvWaitKey(30);

I also added the GrabCutNS namespace, to differentiate the Image class from the rest of the code (that probably has an Image already).
Code is as usual available online in the SVN repo.
Enjoy!
Roy.