git.maemo.org Git - opencv/blob - samples/c/watershed.cpp

   1 #ifdef _CH_
   2 #pragma package <opencv>
   3 #endif
   4
   5 #ifndef _EiC
   6 #include "cv.h"
   7 #include "highgui.h"
   8 #include <stdio.h>
   9 #include <stdlib.h>
  10 #endif
  11
  12 IplImage* marker_mask = 0;
  13 IplImage* markers = 0;
  14 IplImage* img0 = 0, *img = 0, *img_gray = 0, *wshed = 0;
  15 CvPoint prev_pt = {-1,-1};
  16
  17 void on_mouse( int event, int x, int y, int flags, void* param )
  18 {
  19     if( !img )
  20         return;
  21
  22     if( event == CV_EVENT_LBUTTONUP || !(flags & CV_EVENT_FLAG_LBUTTON) )
  23         prev_pt = cvPoint(-1,-1);
  24     else if( event == CV_EVENT_LBUTTONDOWN )
  25         prev_pt = cvPoint(x,y);
  26     else if( event == CV_EVENT_MOUSEMOVE && (flags & CV_EVENT_FLAG_LBUTTON) )
  27     {
  28         CvPoint pt = cvPoint(x,y);
  29         if( prev_pt.x < 0 )
  30             prev_pt = pt;
  31         cvLine( marker_mask, prev_pt, pt, cvScalarAll(255), 5, 8, 0 );
  32         cvLine( img, prev_pt, pt, cvScalarAll(255), 5, 8, 0 );
  33         prev_pt = pt;
  34         cvShowImage( "image", img );
  35     }
  36 }
  37
  38
  39 int main( int argc, char** argv )
  40 {
  41     char* filename = argc >= 2 ? argv[1] : (char*)"fruits.jpg";
  42     CvRNG rng = cvRNG(-1);
  43
  44     if( (img0 = cvLoadImage(filename,1)) == 0 )
  45         return 0;
  46
  47     printf( "Hot keys: \n"
  48             "\tESC - quit the program\n"
  49             "\tr - restore the original image\n"
  50             "\tw or ENTER - run watershed algorithm\n"
  51             "\t\t(before running it, roughly mark the areas on the image)\n"
  52             "\t  (before that, roughly outline several markers on the image)\n" );
  53
  54     cvNamedWindow( "image", 1 );
  55     cvNamedWindow( "watershed transform", 1 );
  56
  57     img = cvCloneImage( img0 );
  58     img_gray = cvCloneImage( img0 );
  59     wshed = cvCloneImage( img0 );
  60     marker_mask = cvCreateImage( cvGetSize(img), 8, 1 );
  61     markers = cvCreateImage( cvGetSize(img), IPL_DEPTH_32S, 1 );
  62     cvCvtColor( img, marker_mask, CV_BGR2GRAY );
  63     cvCvtColor( marker_mask, img_gray, CV_GRAY2BGR );
  64
  65     cvZero( marker_mask );
  66     cvZero( wshed );
  67     cvShowImage( "image", img );
  68     cvShowImage( "watershed transform", wshed );
  69     cvSetMouseCallback( "image", on_mouse, 0 );
  70
  71     for(;;)
  72     {
  73         int c = cvWaitKey(0);
  74
  75         if( (char)c == 27 )
  76             break;
  77
  78         if( (char)c == 'r' )
  79         {
  80             cvZero( marker_mask );
  81             cvCopy( img0, img );
  82             cvShowImage( "image", img );
  83         }
  84
  85         if( (char)c == 'w' || (char)c == '\n' )
  86         {
  87             CvMemStorage* storage = cvCreateMemStorage(0);
  88             CvSeq* contours = 0;
  89             CvMat* color_tab;
  90             int i, j, comp_count = 0;
  91             //cvSaveImage( "wshed_mask.png", marker_mask );
  92             //marker_mask = cvLoadImage( "wshed_mask.png", 0 );
  93             cvFindContours( marker_mask, storage, &contours, sizeof(CvContour),
  94                             CV_RETR_CCOMP, CV_CHAIN_APPROX_SIMPLE );
  95             cvZero( markers );
  96             for( ; contours != 0; contours = contours->h_next, comp_count++ )
  97             {
  98                 cvDrawContours( markers, contours, cvScalarAll(comp_count+1),
  99                                 cvScalarAll(comp_count+1), -1, -1, 8, cvPoint(0,0) );
 100             }
 101
 102             color_tab = cvCreateMat( 1, comp_count, CV_8UC3 );
 103             for( i = 0; i < comp_count; i++ )
 104             {
 105                 uchar* ptr = color_tab->data.ptr + i*3;
 106                 ptr[0] = (uchar)(cvRandInt(&rng)%180 + 50);
 107                 ptr[1] = (uchar)(cvRandInt(&rng)%180 + 50);
 108                 ptr[2] = (uchar)(cvRandInt(&rng)%180 + 50);
 109             }
 110
 111             {
 112             double t = (double)cvGetTickCount();
 113             cvWatershed( img0, markers );
 114             t = (double)cvGetTickCount() - t;
 115             printf( "exec time = %gms\n", t/(cvGetTickFrequency()*1000.) );
 116             }
 117
 118             // paint the watershed image
 119             for( i = 0; i < markers->height; i++ )
 120                 for( j = 0; j < markers->width; j++ )
 121                 {
 122                     int idx = CV_IMAGE_ELEM( markers, int, i, j );
 123                     uchar* dst = &CV_IMAGE_ELEM( wshed, uchar, i, j*3 );
 124                     if( idx == -1 )
 125                         dst[0] = dst[1] = dst[2] = (uchar)255;
 126                     else if( idx <= 0 || idx > comp_count )
 127                         dst[0] = dst[1] = dst[2] = (uchar)0; // should not get here
 128                     else
 129                     {
 130                         uchar* ptr = color_tab->data.ptr + (idx-1)*3;
 131                         dst[0] = ptr[0]; dst[1] = ptr[1]; dst[2] = ptr[2];
 132                     }
 133                 }
 134
 135             cvAddWeighted( wshed, 0.5, img_gray, 0.5, 0, wshed );
 136             cvShowImage( "watershed transform", wshed );
 137             cvReleaseMemStorage( &storage );
 138             cvReleaseMat( &color_tab );
 139         }
 140     }
 141
 142     return 1;
 143 }
 144
 145 #ifdef _EiC
 146 main(1,"watershed.cpp");
 147 #endif