Interactive segmentation of objects based on python OpenCV GrabCut algorithm

20190516

previously on

Find a piece of fun code. Fishing time is on the rise, which improves some bug s and interaction logic of the original code.

Code

import numpy as np
import cv2
from matplotlib import pyplot as plt
import  warnings
warnings.filterwarnings("ignore", module="matplotlib")

imgpath = "E:/20190515-XML_PNG/10.tif_thumb.png"
img = cv2.imread(imgpath)


Coords1x, Coords1y = 'NA', 'NA'
Coords2x, Coords2y = 'NA', 'NA'


def OnClick(event):
    # Gets the position of the mouse when the mouse is "pressed"
    global Coords1x, Coords1y
    if event.button == 1:
        try:
            Coords1x = int(event.xdata)
            Coords1y = int(event.ydata)
        except:
            Coords1x = event.xdata
            Coords1y = event.ydata
        print("#### Upper left coordinate:“ , Coords1x, Coords1y)
 
 
def OnMouseMotion(event):
    # Gets the position of the mouse when it "moves"
    global Coords2x, Coords2y
    if event.button == 3:
        try:
            Coords2x = int(event.xdata)
            Coords2y = int(event.ydata)
        except:
            Coords2x = event.xdata
            Coords2y = event.ydata
        print("#### Lower right coordinate:“ , Coords2x, Coords2x)


def OnMouseRelease(event):
    if event.button == 2:
        fig = plt.gca()
        img = cv2.imread(imgpath)
        # Create a Mask of the same shape as the loaded image
        mask = np.zeros(img.shape[:2], np.uint8)

        # For the array used inside the algorithm, you must create two 0 arrays of type np.float64 with the size of (1, 65)
        bgdModel = np.zeros((1, 65), np.float64)
        fgdModel = np.zeros((1, 65), np.float64)
 
        # Calculate the rectangular area of artificial foreground (rect.x,rect.y,rect.width,rect.height)
        try:
            rect = (Coords1x, Coords1y, Coords2x - Coords1x, Coords2y - Coords1y)
            print('#### Split area: ',rect)
            print('#### It's a little slow... ')
            iterCount = 5
            cv2.grabCut(img, mask, rect, bgdModel, fgdModel, iterCount, cv2.GC_INIT_WITH_RECT)
     
            mask2 = np.where((mask == 2) | (mask == 0), 0, 1).astype('uint8')
            img = img * mask2[:, :, np.newaxis]
            plt.subplot(121), plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
            plt.subplot(122), plt.imshow(cv2.cvtColor(cv2.imread(imgpath), cv2.COLOR_BGR2RGB))
            fig.figure.canvas.draw()
            print('May the force be with me!')
        except:
            print('#### First left then right ")


# Predrawn picture
fig = plt.figure()
plt.subplot(121), plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
plt.subplot(122), plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
plt.colorbar()

# Left mouse button, select the upper left corner of the divided area (rectangle)
fig.canvas.mpl_connect('button_press_event', OnClick)
# Right click to select the lower right corner of the divided area (rectangle)
fig.canvas.mpl_connect('button_press_event', OnMouseMotion)
# Middle click to split the selected area
fig.canvas.mpl_connect('button_press_event', OnMouseRelease)
plt.show()
 
# Connect mouse click events
# fig.canvas.mpl_connect('motion_notify_event', OnMouseMotion)
# Connect mouse move events
# fig.canvas.mpl_connect('button_release_event', OnMouseRelease)
# plt.show()

Operation instructions

  • 1. Code flow
    • Select a rectangular area in the image area,
    • The rectangular area is determined by two points (upper left corner and lower right corner)
    • Use GrabCut algorithm to segment foreground in this region
  • 2. Interaction process
    • Click the left mouse button in the image area to select the upper left corner of the rectangular area
    • Right click the image area to select the lower right corner of the rectangle area
    • Click the middle key to execute the GrabCut algorithm in the selected rectangular area

Display effect

Added by sinista on Sat, 09 Nov 2019 17:46:38 +0200