-
Notifications
You must be signed in to change notification settings - Fork 4
/
grab-cut.py
90 lines (77 loc) · 3.12 KB
/
grab-cut.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
# import the necessary packages
import cv2
import sys
import os
import numpy as np
import argparse
import time
sys.path.append(os.path.abspath('./modules/'))
import detection
import helpers
ap = argparse.ArgumentParser()
ap.add_argument("-d", "--debug", type=str,
default = None,
help="Show each mask detection for Debugging")
args = ap.parse_args()
image_path = helpers.get_random_test_image()
image = cv2.imread(image_path,1)
ball = detection.GolfBallDetection(image)
if ball:
(x,y,w,h) = ball
#The boundaries are often not very accurate
#so we can expand on them a little to make the mask better
x -= 10
y -= 10
w += 30
h += 30
x_plus_w = x+w
y_plus_h = y+h
mask = np.zeros(image.shape[:2], dtype="uint8")
# allocate memory for two arrays that the GrabCut algorithm internally
# used 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()
iter_count = 5
(mask, bgModel, fgModel) = cv2.grabCut(image, mask, ball, bgModel,
fgModel, iterCount=iter_count, 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),
)
if args.debug:
# 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)
cv2.waitKey(0)
# 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)
# Label the boundaries of the grab cut mask
detection.draw_boundaries_and_label(image,(x,y),(w,h),(0,255,0),"Boundaries adjusted")
# show the input image followed by the mask and output generated by
# GrabCut and bitwise masking
cv2.imshow("Input", image)
if args.debug:
cv2.imshow("GrabCut Mask", outputMask)
cv2.imshow("GrabCut Output", output)
cv2.waitKey(0)