-
Notifications
You must be signed in to change notification settings - Fork 0
/
circle_detection.py
74 lines (64 loc) · 2.54 KB
/
circle_detection.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
import numpy as np
import cv2
import scipy.spatial
from ball_classifier import classifyBall
def findCircles():
# read image - 0 is greyscale, 1 - color
table_img = cv2.imread('training_sets/tables/extable6.png', 1)
table_img_col = table_img.copy()
table_img_grey = cv2.cvtColor(table_img, cv2.COLOR_BGR2GRAY)
table_orig = table_img_grey.copy()
# smooth
table_img_grey = cv2.blur(table_img_grey, (3,3))
# perform canny edge detection
table_canny = cv2.Canny(table_img_grey, 15, 30)
t_c_copy = table_canny.copy()
# Perform Hough circle transform
circles = cv2.HoughCircles(table_canny, cv2.HOUGH_GRADIENT, 1, 25, param1=90, param2=30, maxRadius=50, minRadius=14)
avgObjRadius = 0
stripes = []
solids = []
cueBall = (0,0)
pockets = []
if circles is not None:
print("Found circles")
circles = np.round(circles[0, :]).astype("int")
totAvgRadius = sum(i[2] for i in circles) // len(circles)
objBallCounter = 0
for x, y, r in circles:
if r <= totAvgRadius:
objBallCounter += 1
avgObjRadius += r
avgObjRadius = avgObjRadius // objBallCounter
for x, y, r in circles:
if r > 30:
pockets.append([x, y, r])
cv2.circle(table_img, (x, y), r, (0, 210, 30), 3)
else:
# store pixels within circle below
ball = isolateBall(x, y, avgObjRadius, table_img)
ballType = classifyBall(ball)
if ballType == "stripe":
stripes.append((x, y))
elif ballType == "solid":
solids.append((x, y))
elif ballType == "cue":
cueBall = (x, y)
else:
raise Exception("Ball can not be classified. X= " + x + " Y= " + y)
cv2.circle(table_img, (x, y), avgObjRadius, (150, 100, 255), 4)
#concatenate before+after images
img = np.concatenate((table_img_col, cv2.cvtColor(t_c_copy, cv2.COLOR_GRAY2BGR), table_img), axis=0)
filename = 'img.png'
cv2.imwrite(filename, img)
return filename
def isolateBall(x, y, r, img):
height, width = img.shape[:2]
mask = np.zeros((height, width, 3), np.uint8)
height1, width1 = mask.shape[:2]
ball = []
for i in range(x - r, x + r): # columns (width)
for j in range(y - r, y + r): # rows (height)
if (scipy.spatial.distance.euclidean([x,y], [i,j]) < r):
mask[j][i] = img[j][i]
return mask