1793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler#!/usr/bin/env python 2793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler''' 3793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler=============================================================================== 4793ee12c6df9cad3806238d32528c49a3ff9331dNoah PreslerInteractive Image Segmentation using GrabCut algorithm. 5793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 6793ee12c6df9cad3806238d32528c49a3ff9331dNoah PreslerThis sample shows interactive image segmentation using grabcut algorithm. 7793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 8793ee12c6df9cad3806238d32528c49a3ff9331dNoah PreslerUSAGE: 9793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler python grabcut.py <filename> 10793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 11793ee12c6df9cad3806238d32528c49a3ff9331dNoah PreslerREADME FIRST: 12793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler Two windows will show up, one for input and one for output. 13793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 14793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler At first, in input window, draw a rectangle around the object using 15793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslermouse right button. Then press 'n' to segment the object (once or a few times) 16793ee12c6df9cad3806238d32528c49a3ff9331dNoah PreslerFor any finer touch-ups, you can press any of the keys below and draw lines on 17793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerthe areas you want. Then again press 'n' for updating the output. 18793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 19793ee12c6df9cad3806238d32528c49a3ff9331dNoah PreslerKey '0' - To select areas of sure background 20793ee12c6df9cad3806238d32528c49a3ff9331dNoah PreslerKey '1' - To select areas of sure foreground 21793ee12c6df9cad3806238d32528c49a3ff9331dNoah PreslerKey '2' - To select areas of probable background 22793ee12c6df9cad3806238d32528c49a3ff9331dNoah PreslerKey '3' - To select areas of probable foreground 23793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 24793ee12c6df9cad3806238d32528c49a3ff9331dNoah PreslerKey 'n' - To update the segmentation 25793ee12c6df9cad3806238d32528c49a3ff9331dNoah PreslerKey 'r' - To reset the setup 26793ee12c6df9cad3806238d32528c49a3ff9331dNoah PreslerKey 's' - To save the results 27793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler=============================================================================== 28793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler''' 29793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 30793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerimport numpy as np 31793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerimport cv2 32793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerimport sys 33793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 34793ee12c6df9cad3806238d32528c49a3ff9331dNoah PreslerBLUE = [255,0,0] # rectangle color 35793ee12c6df9cad3806238d32528c49a3ff9331dNoah PreslerRED = [0,0,255] # PR BG 36793ee12c6df9cad3806238d32528c49a3ff9331dNoah PreslerGREEN = [0,255,0] # PR FG 37793ee12c6df9cad3806238d32528c49a3ff9331dNoah PreslerBLACK = [0,0,0] # sure BG 38793ee12c6df9cad3806238d32528c49a3ff9331dNoah PreslerWHITE = [255,255,255] # sure FG 39793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 40793ee12c6df9cad3806238d32528c49a3ff9331dNoah PreslerDRAW_BG = {'color' : BLACK, 'val' : 0} 41793ee12c6df9cad3806238d32528c49a3ff9331dNoah PreslerDRAW_FG = {'color' : WHITE, 'val' : 1} 42793ee12c6df9cad3806238d32528c49a3ff9331dNoah PreslerDRAW_PR_FG = {'color' : GREEN, 'val' : 3} 43793ee12c6df9cad3806238d32528c49a3ff9331dNoah PreslerDRAW_PR_BG = {'color' : RED, 'val' : 2} 44793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 45793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler# setting up flags 46793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerrect = (0,0,1,1) 47793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerdrawing = False # flag for drawing curves 48793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerrectangle = False # flag for drawing rect 49793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerrect_over = False # flag to check if rect drawn 50793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerrect_or_mask = 100 # flag for selecting rect or mask mode 51793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslervalue = DRAW_FG # drawing initialized to FG 52793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerthickness = 3 # brush thickness 53793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 54793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerdef onmouse(event,x,y,flags,param): 55793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler global img,img2,drawing,value,mask,rectangle,rect,rect_or_mask,ix,iy,rect_over 56793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 57793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler # Draw Rectangle 58793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler if event == cv2.EVENT_RBUTTONDOWN: 59793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler rectangle = True 60793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler ix,iy = x,y 61793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 62793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler elif event == cv2.EVENT_MOUSEMOVE: 63793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler if rectangle == True: 64793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler img = img2.copy() 65793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler cv2.rectangle(img,(ix,iy),(x,y),BLUE,2) 66793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler rect = (min(ix,x),min(iy,y),abs(ix-x),abs(iy-y)) 67793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler rect_or_mask = 0 68793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 69793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler elif event == cv2.EVENT_RBUTTONUP: 70793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler rectangle = False 71793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler rect_over = True 72793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler cv2.rectangle(img,(ix,iy),(x,y),BLUE,2) 73793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler rect = (min(ix,x),min(iy,y),abs(ix-x),abs(iy-y)) 74793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler rect_or_mask = 0 75793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler print " Now press the key 'n' a few times until no further change \n" 76793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 77793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler # draw touchup curves 78793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 79793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler if event == cv2.EVENT_LBUTTONDOWN: 80793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler if rect_over == False: 81793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler print "first draw rectangle \n" 82793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler else: 83793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler drawing = True 84793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler cv2.circle(img,(x,y),thickness,value['color'],-1) 85793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler cv2.circle(mask,(x,y),thickness,value['val'],-1) 86793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 87793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler elif event == cv2.EVENT_MOUSEMOVE: 88793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler if drawing == True: 89793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler cv2.circle(img,(x,y),thickness,value['color'],-1) 90793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler cv2.circle(mask,(x,y),thickness,value['val'],-1) 91793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 92793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler elif event == cv2.EVENT_LBUTTONUP: 93793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler if drawing == True: 94793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler drawing = False 95793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler cv2.circle(img,(x,y),thickness,value['color'],-1) 96793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler cv2.circle(mask,(x,y),thickness,value['val'],-1) 97793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 98793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler# print documentation 99793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerprint __doc__ 100793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 101793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler# Loading images 102793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerif len(sys.argv) == 2: 103793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler filename = sys.argv[1] # for drawing purposes 104793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerelse: 105793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler print "No input image given, so loading default image, ../data/lena.jpg \n" 106793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler print "Correct Usage: python grabcut.py <filename> \n" 107793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler filename = '../data/lena.jpg' 108793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 109793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerimg = cv2.imread(filename) 110793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerimg2 = img.copy() # a copy of original image 111793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslermask = np.zeros(img.shape[:2],dtype = np.uint8) # mask initialized to PR_BG 112793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presleroutput = np.zeros(img.shape,np.uint8) # output image to be shown 113793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 114793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler# input and output windows 115793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslercv2.namedWindow('output') 116793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslercv2.namedWindow('input') 117793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslercv2.setMouseCallback('input',onmouse) 118793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslercv2.moveWindow('input',img.shape[1]+10,90) 119793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 120793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerprint " Instructions: \n" 121793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerprint " Draw a rectangle around the object using right mouse button \n" 122793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 123793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslerwhile(1): 124793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 125793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler cv2.imshow('output',output) 126793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler cv2.imshow('input',img) 127793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler k = 0xFF & cv2.waitKey(1) 128793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 129793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler # key bindings 130793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler if k == 27: # esc to exit 131793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler break 132793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler elif k == ord('0'): # BG drawing 133793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler print " mark background regions with left mouse button \n" 134793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler value = DRAW_BG 135793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler elif k == ord('1'): # FG drawing 136793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler print " mark foreground regions with left mouse button \n" 137793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler value = DRAW_FG 138793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler elif k == ord('2'): # PR_BG drawing 139793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler value = DRAW_PR_BG 140793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler elif k == ord('3'): # PR_FG drawing 141793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler value = DRAW_PR_FG 142793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler elif k == ord('s'): # save image 143793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler bar = np.zeros((img.shape[0],5,3),np.uint8) 144793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler res = np.hstack((img2,bar,img,bar,output)) 145793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler cv2.imwrite('grabcut_output.png',res) 146793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler print " Result saved as image \n" 147793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler elif k == ord('r'): # reset everything 148793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler print "resetting \n" 149793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler rect = (0,0,1,1) 150793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler drawing = False 151793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler rectangle = False 152793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler rect_or_mask = 100 153793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler rect_over = False 154793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler value = DRAW_FG 155793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler img = img2.copy() 156793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler mask = np.zeros(img.shape[:2],dtype = np.uint8) # mask initialized to PR_BG 157793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler output = np.zeros(img.shape,np.uint8) # output image to be shown 158793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler elif k == ord('n'): # segment the image 159793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler print """ For finer touchups, mark foreground and background after pressing keys 0-3 160793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler and again press 'n' \n""" 161793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler if (rect_or_mask == 0): # grabcut with rect 162793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler bgdmodel = np.zeros((1,65),np.float64) 163793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler fgdmodel = np.zeros((1,65),np.float64) 164793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler cv2.grabCut(img2,mask,rect,bgdmodel,fgdmodel,1,cv2.GC_INIT_WITH_RECT) 165793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler rect_or_mask = 1 166793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler elif rect_or_mask == 1: # grabcut with mask 167793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler bgdmodel = np.zeros((1,65),np.float64) 168793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler fgdmodel = np.zeros((1,65),np.float64) 169793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler cv2.grabCut(img2,mask,rect,bgdmodel,fgdmodel,1,cv2.GC_INIT_WITH_MASK) 170793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 171793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler mask2 = np.where((mask==1) + (mask==3),255,0).astype('uint8') 172793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler output = cv2.bitwise_and(img2,img2,mask=mask2) 173793ee12c6df9cad3806238d32528c49a3ff9331dNoah Presler 174793ee12c6df9cad3806238d32528c49a3ff9331dNoah Preslercv2.destroyAllWindows() 175