vcs.maemo.org Git - opencv/blob - samples/python/watershed.py

   1 #!/usr/bin/python
   2 from opencv.cv import *
   3 from opencv.highgui import *
   4 import sys
   5
   6 marker_mask = None;
   7 markers = None;
   8 img0 = None
   9 img = None
  10 img_gray = None
  11 wshed = None
  12 prev_pt = cvPoint(-1,-1)
  13
  14 def on_mouse( event, x, y, flags, param ):
  15     global prev_pt
  16     if( not img ):
  17         return;
  18     if( event == CV_EVENT_LBUTTONUP or not (flags & CV_EVENT_FLAG_LBUTTON) ):
  19         prev_pt = cvPoint(-1,-1);
  20     elif( event == CV_EVENT_LBUTTONDOWN ):
  21         prev_pt = cvPoint(x,y);
  22     elif( event == CV_EVENT_MOUSEMOVE and (flags & CV_EVENT_FLAG_LBUTTON) ):
  23         pt = cvPoint(x,y);
  24         if( prev_pt.x < 0 ):
  25             prev_pt = pt;
  26         cvLine( marker_mask, prev_pt, pt, cvScalarAll(255), 5, 8, 0 );
  27         cvLine( img, prev_pt, pt, cvScalarAll(255), 5, 8, 0 );
  28         prev_pt = pt;
  29         cvShowImage( "image", img );
  30
  31 if __name__ == "__main__":
  32     filename = "../c/fruits.jpg"
  33     if len(sys.argv)>1:
  34         filename = sys.argv[1]
  35
  36     rng = cvRNG(-1);
  37     img0 = cvLoadImage(filename,1)
  38     if not img0:
  39         print "Error opening image '%s'" % filename
  40         sys.exit(-1)
  41
  42     print "Hot keys:"
  43     print "\tESC - quit the program"
  44     print "\tr - restore the original image"
  45     print "\tw - run watershed algorithm"
  46     print "\t  (before that, roughly outline several markers on the image)"
  47
  48     cvNamedWindow( "image", 1 );
  49     cvNamedWindow( "watershed transform", 1 );
  50
  51     img = cvCloneImage( img0 );
  52     img_gray = cvCloneImage( img0 );
  53     wshed = cvCloneImage( img0 );
  54     marker_mask = cvCreateImage( cvGetSize(img), 8, 1 );
  55     markers = cvCreateImage( cvGetSize(img), IPL_DEPTH_32S, 1 );
  56
  57     cvCvtColor( img, marker_mask, CV_BGR2GRAY );
  58     cvCvtColor( marker_mask, img_gray, CV_GRAY2BGR );
  59
  60     cvZero( marker_mask );
  61     cvZero( wshed );
  62
  63     cvShowImage( "image", img );
  64     cvShowImage( "watershed transform", wshed );
  65
  66     cvSetMouseCallback( "image", on_mouse, None );
  67     while True:
  68         c = cvWaitKey(0);
  69         if c=='\x1b':
  70             break;
  71         if c == 'r':
  72             cvZero( marker_mask );
  73             cvCopy( img0, img );
  74             cvShowImage( "image", img );
  75         if c == 'w':
  76             storage = cvCreateMemStorage(0);
  77             comp_count = 0;
  78             #cvSaveImage( "wshed_mask.png", marker_mask );
  79             #marker_mask = cvLoadImage( "wshed_mask.png", 0 );
  80             nb_cont, contours = cvFindContours( marker_mask, storage, sizeof_CvContour,
  81                             CV_RETR_CCOMP, CV_CHAIN_APPROX_SIMPLE );
  82             cvZero( markers );
  83             while contours:
  84                 cvDrawContours( markers, contours, cvScalarAll(comp_count+1),
  85                                 cvScalarAll(comp_count+1), -1, -1, 8, cvPoint(0,0) );
  86                 contours=contours.h_next
  87                 comp_count+=1
  88             color_tab = cvCreateMat( comp_count, 1, CV_8UC3 );
  89             for i in range(comp_count):
  90                 color_tab[i] = cvScalar( cvRandInt(rng)%180 + 50,
  91                                  cvRandInt(rng)%180 + 50,
  92                                  cvRandInt(rng)%180 + 50 );
  93             t = cvGetTickCount();
  94             cvWatershed( img0, markers );
  95             t = cvGetTickCount() - t;
  96             #print "exec time = %f" % t/(cvGetTickFrequency()*1000.)
  97
  98             cvSet( wshed, cvScalarAll(255) );
  99
 100             # paint the watershed image
 101             for j in range(markers.height):
 102                 for i in range(markers.width):
 103                     idx = markers[j,i]
 104                     if idx==-1:
 105                         continue
 106                     idx = idx-1
 107                     wshed[j,i] = color_tab[idx,0]
 108
 109             cvAddWeighted( wshed, 0.5, img_gray, 0.5, 0, wshed );
 110             cvShowImage( "watershed transform", wshed );
 111             cvWaitKey();