Skip to content
Permalink
Branch: master
Find file Copy path
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
186 lines (162 sloc) 6.97 KB
#!/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
===============================================================================
'''
# Python 2/3 compatibility
from __future__ import print_function
import numpy as np
import cv2 as cv
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,originalimg,drawing,value,mask,rectangle,rect,rect_or_mask,ix,iy,rect_over
# Draw Rectangle
if event == cv.EVENT_RBUTTONDOWN:
rectangle = True
ix,iy = x,y
elif event == cv.EVENT_MOUSEMOVE:
if rectangle == True:
img = img2.copy()
cv.rectangle(img,(ix,iy),(x,y),BLUE,2)
rect = (min(ix,x),min(iy,y),abs(ix-x),abs(iy-y))
rect_or_mask = 0
elif event == cv.EVENT_RBUTTONUP:
rectangle = False
rect_over = True
cv.rectangle(img,(ix,iy),(x,y),BLUE,2)
rect = (min(ix,x),min(iy,y),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 == cv.EVENT_LBUTTONDOWN:
if rect_over == False:
print("first draw rectangle \n")
else:
drawing = True
cv.circle(img,(x,y),thickness,value['color'],-1)
cv.circle(mask,(x,y),thickness,value['val'],-1)
elif event == cv.EVENT_MOUSEMOVE:
if drawing == True:
cv.circle(img,(x,y),thickness,value['color'],-1)
cv.circle(mask,(x,y),thickness,value['val'],-1)
elif event == cv.EVENT_LBUTTONUP:
if drawing == True:
drawing = False
cv.circle(img,(x,y),thickness,value['color'],-1)
cv.circle(mask,(x,y),thickness,value['val'],-1)
if __name__ == '__main__':
# 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, ../data/lena.jpg \n")
print("Correct Usage: python grabcut.py <filename> \n")
filename = '../data/lena.jpg'
img = cv.imread(filename)
img2 = img.copy() # a copy of original image
originalimg = 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
# output.fill(255)
output[:] = 255
# input and output windows
cv.namedWindow('output')
cv.namedWindow('input')
cv.setMouseCallback('input',onmouse)
cv.moveWindow('input',img.shape[1]+10,90)
print(" Instructions: \n")
print(" Draw a rectangle around the object using right mouse button \n")
while(1):
cv.imshow('output',output)
cv.imshow('input',img)
k = cv.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))
shortfilename = filename.split("/")[-1]
newshortfilename = 'grabcut_' + shortfilename
newfilepath = filename.replace(shortfilename,newshortfilename)
# save the new img in the same path with the original img
cv.imwrite(newfilepath,output)
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
output0 = 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)
cv.grabCut(img2,mask,rect,bgdmodel,fgdmodel,1,cv.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)
cv.grabCut(img2,mask,rect,bgdmodel,fgdmodel,1,cv.GC_INIT_WITH_MASK)
mask2 = np.where((mask==1) + (mask==3),255,0).astype('uint8')
output = cv.bitwise_and(img2,img2,mask=mask2)
# output0 = cv.bitwise_and(img2,img2,mask=mask2)
# # Change the backgrounf to white
# background = output0 - originalimg
# background[np.where((background > [0,0,0]).all(axis = 2))] =[255,255,255]
# # Add the background and the image
# output = background + output0
cv.destroyAllWindows()
You can’t perform that action at this time.