Update to 2.0.0 tree from current Fremantle build

[opencv] / samples / swig_python / watershed.py

diff --git a/samples/swig_python/watershed.py b/samples/swig_python/watershed.py
new file mode 100755 (executable)
index 0000000..4da9379
--- /dev/null
+++ b/samples/swig_python/watershed.py
@@ -0,0 +1,111 @@
+#!/usr/bin/python
+from opencv.cv import *
+from opencv.highgui import *
+import sys
+
+marker_mask = None;
+markers = None;
+img0 = None
+img = None
+img_gray = None 
+wshed = None
+prev_pt = cvPoint(-1,-1)
+
+def on_mouse( event, x, y, flags, param ):
+    global prev_pt
+    if( not img ):
+        return;
+    if( event == CV_EVENT_LBUTTONUP or not (flags & CV_EVENT_FLAG_LBUTTON) ):
+        prev_pt = cvPoint(-1,-1);
+    elif( event == CV_EVENT_LBUTTONDOWN ):
+        prev_pt = cvPoint(x,y);
+    elif( event == CV_EVENT_MOUSEMOVE and (flags & CV_EVENT_FLAG_LBUTTON) ):
+        pt = cvPoint(x,y);
+        if( prev_pt.x < 0 ):
+            prev_pt = pt;
+        cvLine( marker_mask, prev_pt, pt, cvScalarAll(255), 5, 8, 0 );
+        cvLine( img, prev_pt, pt, cvScalarAll(255), 5, 8, 0 );
+        prev_pt = pt;
+        cvShowImage( "image", img );
+
+if __name__ == "__main__":
+    filename = "../c/fruits.jpg"
+    if len(sys.argv)>1:
+        filename = sys.argv[1]
+
+    rng = cvRNG(-1);
+    img0 = cvLoadImage(filename,1)
+    if not img0:
+        print "Error opening image '%s'" % filename
+        sys.exit(-1)
+
+    print "Hot keys:"
+    print "\tESC - quit the program"
+    print "\tr - restore the original image"
+    print "\tw - run watershed algorithm"
+    print "\t  (before that, roughly outline several markers on the image)"
+
+    cvNamedWindow( "image", 1 );
+    cvNamedWindow( "watershed transform", 1 );
+
+    img = cvCloneImage( img0 );
+    img_gray = cvCloneImage( img0 );
+    wshed = cvCloneImage( img0 );
+    marker_mask = cvCreateImage( cvGetSize(img), 8, 1 );
+    markers = cvCreateImage( cvGetSize(img), IPL_DEPTH_32S, 1 );
+
+    cvCvtColor( img, marker_mask, CV_BGR2GRAY );
+    cvCvtColor( marker_mask, img_gray, CV_GRAY2BGR );
+
+    cvZero( marker_mask );
+    cvZero( wshed );
+
+    cvShowImage( "image", img );
+    cvShowImage( "watershed transform", wshed );
+
+    cvSetMouseCallback( "image", on_mouse, None );
+    while True:
+        c = cvWaitKey(0);
+        if c=='\x1b':
+            break;
+        if c == 'r':
+            cvZero( marker_mask );
+            cvCopy( img0, img );
+            cvShowImage( "image", img );
+        if c == 'w':
+            storage = cvCreateMemStorage(0);
+            comp_count = 0;
+            #cvSaveImage( "wshed_mask.png", marker_mask );
+            #marker_mask = cvLoadImage( "wshed_mask.png", 0 );
+            nb_cont, contours = cvFindContours( marker_mask, storage, sizeof_CvContour,
+                            CV_RETR_CCOMP, CV_CHAIN_APPROX_SIMPLE );
+            cvZero( markers );
+            while contours:
+                cvDrawContours( markers, contours, cvScalarAll(comp_count+1),
+                                cvScalarAll(comp_count+1), -1, -1, 8, cvPoint(0,0) );
+                contours=contours.h_next
+                comp_count+=1
+            color_tab = cvCreateMat( comp_count, 1, CV_8UC3 );
+            for i in range(comp_count):
+                color_tab[i] = cvScalar( cvRandInt(rng)%180 + 50, 
+                                 cvRandInt(rng)%180 + 50, 
+                                 cvRandInt(rng)%180 + 50 );
+            t = cvGetTickCount();
+            cvWatershed( img0, markers );
+            t = cvGetTickCount() - t;
+            #print "exec time = %f" % t/(cvGetTickFrequency()*1000.)
+
+            cvSet( wshed, cvScalarAll(255) );
+
+            # paint the watershed image
+            for j in range(markers.height):
+                for i in range(markers.width):
+                    idx = markers[j,i]
+                    if idx==-1:
+                        continue
+                    idx = idx-1
+                    wshed[j,i] = color_tab[idx,0]
+
+            cvAddWeighted( wshed, 0.5, img_gray, 0.5, 0, wshed );
+            cvShowImage( "watershed transform", wshed );
+            cvWaitKey();