From 98960bf201f6ec51f28dfda067d670bb9f85c31b Mon Sep 17 00:00:00 2001 From: abidrahmank Date: Sat, 25 May 2013 00:55:31 +0530 Subject: [PATCH] A new python sample on grabcut --- .../imgproc/doc/miscellaneous_transformations.rst | 2 +- samples/python2/grabcut.py | 174 +++++++++++++++++++++ 2 files changed, 175 insertions(+), 1 deletion(-) create mode 100644 samples/python2/grabcut.py diff --git a/modules/imgproc/doc/miscellaneous_transformations.rst b/modules/imgproc/doc/miscellaneous_transformations.rst index 4ebf6d5..84f9481 100644 --- a/modules/imgproc/doc/miscellaneous_transformations.rst +++ b/modules/imgproc/doc/miscellaneous_transformations.rst @@ -765,7 +765,7 @@ Runs the GrabCut algorithm. * **GC_PR_BGD** defines a possible background pixel. - * **GC_PR_BGD** defines a possible foreground pixel. + * **GC_PR_FGD** defines a possible foreground pixel. :param rect: ROI containing a segmented object. The pixels outside of the ROI are marked as "obvious background". The parameter is only used when ``mode==GC_INIT_WITH_RECT`` . diff --git a/samples/python2/grabcut.py b/samples/python2/grabcut.py new file mode 100644 index 0000000..4de4f39 --- /dev/null +++ b/samples/python2/grabcut.py @@ -0,0 +1,174 @@ +#!/usr/bin/env python +''' +=============================================================================== +Interactive Image Segmentation using GrabCut algorithm. + +This sample shows interactive image segmentation using grabcut algorithm. + +USAGE : + python grabcut.py + +README FIRST: + Two windows will show up, one for input and one for output. + + At first, in input window, draw a rectangle around the object using +mouse right button. Then press 'n' to segment the object (once or a few times) +For any finer touch-ups, you can press any of the keys below and draw lines on +the areas you want. Then again press 'n' for updating the output. + +Key '0' - To select areas of sure background +Key '1' - To select areas of sure foreground +Key '2' - To select areas of probable background +Key '3' - To select areas of probable foreground + +Key 'n' - To update the segmentation +Key 'r' - To reset the setup +Key 's' - To save the results +=============================================================================== +''' + +import numpy as np +import cv2 +import sys + +BLUE = [255,0,0] # rectangle color +RED = [0,0,255] # PR BG +GREEN = [0,255,0] # PR FG +BLACK = [0,0,0] # sure BG +WHITE = [255,255,255] # sure FG + +DRAW_BG = {'color' : BLACK, 'val' : 0} +DRAW_FG = {'color' : WHITE, 'val' : 1} +DRAW_PR_FG = {'color' : GREEN, 'val' : 3} +DRAW_PR_BG = {'color' : RED, 'val' : 2} + +# setting up flags +rect = (0,0,1,1) +drawing = False # flag for drawing curves +rectangle = False # flag for drawing rect +rect_over = False # flag to check if rect drawn +rect_or_mask = 100 # flag for selecting rect or mask mode +value = DRAW_FG # drawing initialized to FG +thickness = 3 # brush thickness + +def onmouse(event,x,y,flags,param): + global img,img2,drawing,value,mask,rectangle,rect,rect_or_mask,ix,iy,rect_over + + # Draw Rectangle + if event == cv2.EVENT_RBUTTONDOWN: + rectangle = True + ix,iy = x,y + + elif event == cv2.EVENT_MOUSEMOVE: + if rectangle == True: + img = img2.copy() + cv2.rectangle(img,(ix,iy),(x,y),BLUE,2) + rect = (ix,iy,abs(ix-x),abs(iy-y)) + rect_or_mask = 0 + + elif event == cv2.EVENT_RBUTTONUP: + rectangle = False + rect_over = True + cv2.rectangle(img,(ix,iy),(x,y),BLUE,2) + rect = (ix,iy,abs(ix-x),abs(iy-y)) + rect_or_mask = 0 + print " Now press the key 'n' a few times until no further change \n" + + # draw touchup curves + + if event == cv2.EVENT_LBUTTONDOWN: + if rect_over == False: + print "first draw rectangle \n" + else: + drawing = True + cv2.circle(img,(x,y),thickness,value['color'],-1) + cv2.circle(mask,(x,y),thickness,value['val'],-1) + + elif event == cv2.EVENT_MOUSEMOVE: + if drawing == True: + cv2.circle(img,(x,y),thickness,value['color'],-1) + cv2.circle(mask,(x,y),thickness,value['val'],-1) + + elif event == cv2.EVENT_LBUTTONUP: + if drawing == True: + drawing = False + cv2.circle(img,(x,y),thickness,value['color'],-1) + cv2.circle(mask,(x,y),thickness,value['val'],-1) + +# print documentation +print __doc__ + +# Loading images +if len(sys.argv) == 2: + filename = sys.argv[1] # for drawing purposes +else: + print "No input image given, so loading default image, lena.jpg \n" + print "Correct Usage : python grabcut.py \n" + filename = '../cpp/lena.jpg' + +img = cv2.imread(filename) +img2 = img.copy() # a copy of original image +mask = np.zeros(img.shape[:2],dtype = np.uint8) # mask initialized to PR_BG +output = np.zeros(img.shape,np.uint8) # output image to be shown + +# input and output windows +cv2.namedWindow('output') +cv2.namedWindow('input') +cv2.setMouseCallback('input',onmouse) +cv2.moveWindow('input',img.shape[1]+10,90) + +print " Instructions : \n" +print " Draw a rectangle around the object using right mouse button \n" + +while(1): + + cv2.imshow('output',output) + cv2.imshow('input',img) + k = cv2.waitKey(1) + + # key bindings + if k == 27: # esc to exit + break + elif k == ord('0'): # BG drawing + print " mark background regions with left mouse button \n" + value = DRAW_BG + elif k == ord('1'): # FG drawing + print " mark foreground regions with left mouse button \n" + value = DRAW_FG + elif k == ord('2'): # PR_BG drawing + value = DRAW_PR_BG + elif k == ord('3'): # PR_FG drawing + value = DRAW_PR_FG + elif k == ord('s'): # save image + bar = np.zeros((img.shape[0],5,3),np.uint8) + res = np.hstack((img2,bar,img,bar,output)) + cv2.imwrite('grabcut_output.png',res) + print " Result saved as image \n" + elif k == ord('r'): # reset everything + print "resetting \n" + rect = (0,0,1,1) + drawing = False + rectangle = False + rect_or_mask = 100 + rect_over = False + value = DRAW_FG + img = img2.copy() + mask = np.zeros(img.shape[:2],dtype = np.uint8) # mask initialized to PR_BG + output = np.zeros(img.shape,np.uint8) # output image to be shown + elif k == ord('n'): # segment the image + print """ For finer touchups, mark foreground and background after pressing keys 0-3 + and again press 'n' \n""" + if (rect_or_mask == 0): # grabcut with rect + bgdmodel = np.zeros((1,65),np.float64) + fgdmodel = np.zeros((1,65),np.float64) + cv2.grabCut(img2,mask,rect,bgdmodel,fgdmodel,1,cv2.GC_INIT_WITH_RECT) + rect_or_mask = 1 + elif rect_or_mask == 1: # grabcut with mask + bgdmodel = np.zeros((1,65),np.float64) + fgdmodel = np.zeros((1,65),np.float64) + cv2.grabCut(img2,mask,rect,bgdmodel,fgdmodel,1,cv2.GC_INIT_WITH_MASK) + + mask2 = np.where((mask==1) + (mask==3),255,0).astype('uint8') + output = cv2.bitwise_and(img2,img2,mask=mask2) + +cv2.destroyAllWindows() -- 2.7.4