import numpy as np
import cv2
#from matplotlib import pyplot as plt

img = cv2.imread('messi.jpg')
mask = np.zeros(img.shape[:2],np.uint8)

bgdModel = np.zeros((1,65),np.float64)
fgdModel = np.zeros((1,65),np.float64)

rect = (50,50,380,275)
#cv2.grabCut(img,mask,rect,bgdModel,fgdModel,5,cv2.GC_INIT_WITH_RECT)

#mask2 = np.where((mask==2)|(mask==0),0,1).astype('uint8')
#img = img*mask2[:,:,np.newaxis]

cv2.rectangle(img,(50,50),(380,275),(0,255,0),3)

# newmask is the mask image I manually labelled
newmask = cv2.imread('marcasFG-BG.jpg',0)
# whereever it is marked white (sure foreground), change mask=1
# whereever it is marked black (sure background), change mask=0
mask[newmask == 0] = 0
mask[newmask == 255] = 1

cv2.imshow("original",img)
print len(mask)
res = cv2.grabCut(img,mask,None,bgdModel,fgdModel,15,cv2.GC_INIT_WITH_MASK)

cv2.imshow("desGrabCut",img)
print res
mask = np.where((mask==2)|(mask==0),0,1).astype('uint8')
img = img*mask[:,:,np.newaxis]

cv2.imshow("resultado",img)
cv2.imshow("mask",newmask)
k=cv2.waitKey(0)&0xFF#
if k== 27: #waitforESCkeytoexit
        cv2.destroyAllWindows()
elif k == ord('s'):
        cv2.imwrite('messisegmentada1.jpg',img)
        cv2.destroyAllWindows()