A new python sample on grabcut

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`` .
 

samples/python2/grabcut.py

+#!/usr/bin/env python
+'''
+===============================================================================
+Interactive Image Segmentation using GrabCut algorithm.
+
+This sample shows interactive image segmentation using grabcut algorithm.
+
+USAGE :
+    python grabcut.py <filename>
+
+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 <filename> \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()