Grabcut image segmentation in opencv

1.GrabCut

Get foreground objects through interaction ;
principle ：
（1） The user specifies the general area of the foreground , The rest is the background area ;
（2） Users can also specify certain places as foreground or background ;
（3）GrabCut The model tree is formed by analyzing the foreground objects in the way of subsection iteration ;
（4） Determine whether a pixel is foreground or background by weight ;

 grabCut(img, mask, rect, bgdModel, fgdModel, iterCount, mode=None): 

Img: Input the original image ;
Mask: The mask generated after segmentation ;
Rect: The size of the area selected with the mouse ;
bgdModel：np.float64 ; The size is 1-65 Array with zero ;
fgdModel：np.float64 ; The size is 1-65 Array with zero ;
Intercount: The number of iterations ;
Mode:GC_INIT_WITH_RECT;GC_INIT_WITH_MASK;

Output Mask Corresponding to each pixel value ：
BGD： background -0;
FGD： prospects -1;
PR_BGD： Maybe the background -2;
PR_FGD： It could be the future -3;

import os
import cv2
import numpy as np

# Define structure classes 
class App:
    flag=False
    startX=0
    startY=0
    rect=(0,0,0,0)
    def onmouse(self,event,x,y,flags,param):
        if(event&cv2.EVENT_LBUTTONDOWN==cv2.EVENT_LBUTTONDOWN):
            self.flag=True
            self.startX=x
            self.startY=y
            print(' Press the left key ')
        elif(event&cv2.EVENT_LBUTTONUP==cv2.EVENT_LBUTTONUP):
            self.flag=False
            cv2.rectangle(self.img,pt1=(self.startX,self.startX),pt2=(x,y),color=(0,255,0),thickness=3)
            # structure rect
            self.rect=(min(self.startX,x),min(self.startY,y),abs(self.startX-x),abs(self.startY-y))
            print(' Pop up left button ')
        elif(event&cv2.EVENT_MOUSEMOVE==cv2.EVENT_MOUSEMOVE):
            if self.flag==True:
                # Each move is drawn in the first picture 
                self.img=self.img2.copy()
                cv2.rectangle(self.img, pt1=(self.startX, self.startX), pt2=(x, y), color=(0, 255, 0), thickness=3)
            print(' Mouse movement ')


    def run(self):
        cv2.namedWindow(winname='input',flags=cv2.WINDOW_AUTOSIZE)
        cv2.setMouseCallback('input',self.onmouse)

        self.img=cv2.imread('images/hometomn1.jpg')
        self.img=cv2.resize(src=self.img,dsize=(450,450))
        self.img2=self.img.copy()
        # structure mask
        self.mask=np.zeros(shape=(self.img.shape[0],self.img.shape[1]),dtype=np.uint8)
        self.output=np.zeros(shape=(self.img.shape[0],self.img.shape[1],self.img.shape[2]),dtype=np.uint8)
        while True:
            cv2.imshow('input',self.img)
            cv2.imshow('output',self.output)
            key=cv2.waitKey(100)
            if (key&0xFF==27):
                break

            if key==ord('s'):
                bgmodel=np.zeros((1,65),dtype=np.float64)
                cv2.grabCut(img=self.img2,mask=self.mask,rect=self.rect,bgdModel=bgmodel,
                            fgdModel=bgmodel,iterCount=1,mode=cv2.GC_INIT_WITH_RECT)
            # When judged as foreground , by 255, Otherwise 0
            mask2=np.where((self.mask==1)|(self.mask==3),255,0).astype('uint8')
            self.output=cv2.bitwise_and(src1=self.img2,src2=self.img2,mask=mask2)


        cv2.destroyAllWindows()



if __name__ == '__main__':
    print('Pycharm')
    app=App()
    print(app.run())

demonstration