112 lines
3.7 KiB
Python
112 lines
3.7 KiB
Python
|
#!/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();
|