Wednesday, 28 April 2021

opencv 37 grabcut


try to extract tiger from the scene

grabcut mask is generated

tiger is extracted
#main.py
import numpy as np
import argparse
import time
import cv2
import os

# construct the argument parser and parse the arguments
ap = argparse.ArgumentParser()
ap.add_argument("-i", "--image", type=str,
                default=os.path.sep.join(["images", "adrian.jpg"]),
                help="path to input image that we'll apply GrabCut to")
ap.add_argument("-c", "--iter", type=int, default=10,
                help="# of GrabCut iterations (larger value => slower runtime)")
args = vars(ap.parse_args())

# load the input image from disk and then allocate memory for the
# output mask generated by GrabCut -- this mask should hae the same
# spatial dimensions as the input image
image = cv2.imread(args["image"])

xl, yl, xr, yr = (0, 0, 0, 0)
btn_down = False
img = image.copy()


def drag_event(event, x, y, flags, param):
    global btn_down, xl, yl, xr, yr, img

    if event == cv2.EVENT_LBUTTONUP and btn_down:
        btn_down = False

        cv2.rectangle(img, (xl - 2, yl - 2), (xr + 2, yr + 2), (0, 0, 255), 2)
        cv2.imshow('original', img)

        grab_cut((xl, yl, xr, yr))


    elif event == cv2.EVENT_MOUSEMOVE and btn_down:
        xr, yr = (x, y)
        cv2.rectangle(img, (xl - 2, yl - 2), (xr + 2, yr + 2), (0, 0, 255), 2)
        cv2.imshow('original', img)
        img = image.copy()

    elif event == cv2.EVENT_LBUTTONDOWN:
        btn_down = True
        xl, yl = (x, y)


def grab_cut(rect):
    mask = np.zeros(image.shape[:2], dtype="uint8")

    # allocate memory for two arrays that the GrabCut algorithm internally
    # uses when segmenting the foreground from the background
    fgModel = np.zeros((1, 65), dtype="float")
    bgModel = np.zeros((1, 65), dtype="float")

    # apply GrabCut using the the bounding box segmentation method
    start = time.time()
    (mask, bgModel, fgModel) = cv2.grabCut(image, mask, rect, bgModel,
                                           fgModel, iterCount=args["iter"], mode=cv2.GC_INIT_WITH_RECT)
    end = time.time()
    print("[INFO] applying GrabCut took {:.2f} seconds".format(end - start))

    # the output mask has for possible output values, marking each pixel
    # in the mask as (1) definite background, (2) definite foreground,
    # (3) probable background, and (4) probable foreground
    values = (
        ("Definite Background", cv2.GC_BGD),
        ("Probable Background", cv2.GC_PR_BGD),
        ("Definite Foreground", cv2.GC_FGD),
        ("Probable Foreground", cv2.GC_PR_FGD),
    )

    # loop over the possible GrabCut mask values
    for (name, value) in values:
        # construct a mask that for the current value
        print("[INFO] showing mask for '{}'".format(name))
        valueMask = (mask == value).astype("uint8") * 255

        # display the mask so we can visualize it
        cv2.imshow(name, valueMask)

    # we'll set all definite background and probable background pixels
    # to 0 while definite foreground and probable foreground pixels are
    # set to 1
    outputMask = np.where((mask == cv2.GC_BGD) | (mask == cv2.GC_PR_BGD),
                          0, 1)

    # scale the mask from the range [0, 1] to [0, 255]
    outputMask = (outputMask * 255).astype("uint8")

    # apply a bitwise AND to the image using our mask generated by
    # GrabCut to generate our final output image
    output = cv2.bitwise_and(image, image, mask=outputMask)

    # show the input image followed by the mask and output generated by
    # GrabCut and bitwise masking
    cv2.imshow("Input", image)
    cv2.imshow("GrabCut Mask", outputMask)
    cv2.imshow("GrabCut Output", output)


cv2.imshow("original image", image)
cv2.setMouseCallback('original image', drag_event)

cv2.waitKey(0)

--------------------------
#logs
(venv) C:\Users\zchen\PycharmProjects\opencv>python grabcut.py -i assets/tiger.jpg
[INFO] applying GrabCut took 2.90 seconds
[INFO] showing mask for 'Definite Background'
[INFO] showing mask for 'Probable Background'
[INFO] showing mask for 'Definite Foreground'
[INFO] showing mask for 'Probable Foreground'

reference:

mouse drag event


No comments:

Post a Comment