Home | History | Annotate | Download | only in cpp 
      1 #include <opencv2/core/utility.hpp>
      2 #include "opencv2/imgproc.hpp"
      3 #include "opencv2/imgcodecs.hpp"
      4 #include "opencv2/highgui.hpp"
      5 
      6 #include <cstdio>
      7 #include <iostream>
      8 
      9 using namespace cv;
     10 using namespace std;
     11 
     12 static void help()
     13 {
     14     cout << "\nThis program demonstrates the famous watershed segmentation algorithm in OpenCV: watershed()\n"
     15             "Usage:\n"
     16             "./watershed [image_name -- default is ../data/fruits.jpg]\n" << endl;
     17 
     18 
     19     cout << "Hot keys: \n"
     20         "\tESC - quit the program\n"
     21         "\tr - restore the original image\n"
     22         "\tw or SPACE - run watershed segmentation algorithm\n"
     23         "\t\t(before running it, *roughly* mark the areas to segment on the image)\n"
     24         "\t  (before that, roughly outline several markers on the image)\n";
     25 }
     26 Mat markerMask, img;
     27 Point prevPt(-1, -1);
     28 
     29 static void onMouse( int event, int x, int y, int flags, void* )
     30 {
     31     if( x < 0 || x >= img.cols || y < 0 || y >= img.rows )
     32         return;
     33     if( event == EVENT_LBUTTONUP || !(flags & EVENT_FLAG_LBUTTON) )
     34         prevPt = Point(-1,-1);
     35     else if( event == EVENT_LBUTTONDOWN )
     36         prevPt = Point(x,y);
     37     else if( event == EVENT_MOUSEMOVE && (flags & EVENT_FLAG_LBUTTON) )
     38     {
     39         Point pt(x, y);
     40         if( prevPt.x < 0 )
     41             prevPt = pt;
     42         line( markerMask, prevPt, pt, Scalar::all(255), 5, 8, 0 );
     43         line( img, prevPt, pt, Scalar::all(255), 5, 8, 0 );
     44         prevPt = pt;
     45         imshow("image", img);
     46     }
     47 }
     48 
     49 int main( int argc, char** argv )
     50 {
     51     char* filename = argc >= 2 ? argv[1] : (char*)"../data/fruits.jpg";
     52     Mat img0 = imread(filename, 1), imgGray;
     53 
     54     if( img0.empty() )
     55     {
     56         cout << "Couldn'g open image " << filename << ". Usage: watershed <image_name>\n";
     57         return 0;
     58     }
     59     help();
     60     namedWindow( "image", 1 );
     61 
     62     img0.copyTo(img);
     63     cvtColor(img, markerMask, COLOR_BGR2GRAY);
     64     cvtColor(markerMask, imgGray, COLOR_GRAY2BGR);
     65     markerMask = Scalar::all(0);
     66     imshow( "image", img );
     67     setMouseCallback( "image", onMouse, 0 );
     68 
     69     for(;;)
     70     {
     71         int c = waitKey(0);
     72 
     73         if( (char)c == 27 )
     74             break;
     75 
     76         if( (char)c == 'r' )
     77         {
     78             markerMask = Scalar::all(0);
     79             img0.copyTo(img);
     80             imshow( "image", img );
     81         }
     82 
     83         if( (char)c == 'w' || (char)c == ' ' )
     84         {
     85             int i, j, compCount = 0;
     86             vector<vector<Point> > contours;
     87             vector<Vec4i> hierarchy;
     88 
     89             findContours(markerMask, contours, hierarchy, RETR_CCOMP, CHAIN_APPROX_SIMPLE);
     90 
     91             if( contours.empty() )
     92                 continue;
     93             Mat markers(markerMask.size(), CV_32S);
     94             markers = Scalar::all(0);
     95             int idx = 0;
     96             for( ; idx >= 0; idx = hierarchy[idx][0], compCount++ )
     97                 drawContours(markers, contours, idx, Scalar::all(compCount+1), -1, 8, hierarchy, INT_MAX);
     98 
     99             if( compCount == 0 )
    100                 continue;
    101 
    102             vector<Vec3b> colorTab;
    103             for( i = 0; i < compCount; i++ )
    104             {
    105                 int b = theRNG().uniform(0, 255);
    106                 int g = theRNG().uniform(0, 255);
    107                 int r = theRNG().uniform(0, 255);
    108 
    109                 colorTab.push_back(Vec3b((uchar)b, (uchar)g, (uchar)r));
    110             }
    111 
    112             double t = (double)getTickCount();
    113             watershed( img0, markers );
    114             t = (double)getTickCount() - t;
    115             printf( "execution time = %gms\n", t*1000./getTickFrequency() );
    116 
    117             Mat wshed(markers.size(), CV_8UC3);
    118 
    119             // paint the watershed image
    120             for( i = 0; i < markers.rows; i++ )
    121                 for( j = 0; j < markers.cols; j++ )
    122                 {
    123                     int index = markers.at<int>(i,j);
    124                     if( index == -1 )
    125                         wshed.at<Vec3b>(i,j) = Vec3b(255,255,255);
    126                     else if( index <= 0 || index > compCount )
    127                         wshed.at<Vec3b>(i,j) = Vec3b(0,0,0);
    128                     else
    129                         wshed.at<Vec3b>(i,j) = colorTab[index - 1];
    130                 }
    131 
    132             wshed = wshed*0.5 + imgGray*0.5;
    133             imshow( "watershed transform", wshed );
    134         }
    135     }
    136 
    137     return 0;
    138 }
    139