def _detect_CO_balls(img, balls, cue, display_list):
PARA1 = int(float(img.shape[1]) / 640 * 15 + 0.5)
PARA2 = int(float(img.shape[1]) / 640 * 40 + 0.5)
p_cue_top, p_cue_bottom, cue_length = cue
## determine which ball is cue ball
min_dist2cue_top = img.shape[0] + 1
cue_ball_idx = -1
for idx, ball in enumerate(balls):
center, radius = ball
dist2cue = zc.calc_triangle_area(center, p_cue_top,
p_cue_bottom) * 2 / cue_length
if dist2cue < PARA1:
dist2cue_top = zc.euc_dist(center, p_cue_top)
if dist2cue_top < min_dist2cue_top:
min_dist2cue_top = dist2cue_top
cue_ball_idx = idx
if cue_ball_idx == -1:
rtn_msg = {'status': 'fail', 'message': 'Cannot find cue ball'}
return (rtn_msg, None)
cue_ball = balls[cue_ball_idx]
## determine which ball is object ball
min_dist2cue_top = img.shape[0] + 1
object_ball_idx = -1
for idx, ball in enumerate(balls):
if idx == cue_ball_idx:
continue
center, radius = ball
dist2cue = zc.calc_triangle_area(center, p_cue_top,
p_cue_bottom) * 2 / cue_length
if dist2cue < PARA2:
dist2cue_top = zc.euc_dist(center, p_cue_top)
if dist2cue_top < min_dist2cue_top:
min_dist2cue_top = dist2cue_top
object_ball_idx = idx
if object_ball_idx == -1:
rtn_msg = {'status': 'fail', 'message': 'Cannot find object ball'}
return (rtn_msg, None)
object_ball = balls[object_ball_idx]
if 'CO_balls' in display_list:
img_balls = img.copy()
cv2.circle(img_balls, (int(cue_ball[0][0]), int(cue_ball[0][1])),
int(cue_ball[1]), (255, 255, 255), -1)
cv2.circle(img_balls, (int(object_ball[0][0]), int(object_ball[0][1])),
int(object_ball[1]), (0, 0, 255), -1)
zc.display_image("CO_balls",
img_balls,
resize_max=config.DISPLAY_MAX_PIXEL,
wait_time=config.DISPLAY_WAIT_TIME)
rtn_msg = {'status': 'success'}
return (rtn_msg, (cue_ball, object_ball))
def _detect_aim_point(cue, CO_balls, pocket):
def _angle2fraction(angle):
'''
Calculcate desired fraction of overlap from angle according to fraction aiming system
Reference: http://billiards.colostate.edu/threads/aiming.html, https://youtu.be/2kuJTwQ1M9k
'''
if angle > 48.6: # very thin
return (90 - angle) / 41.4 * 0.25
elif angle > 30: # thin
return (48.6 - angle) / 28.6 * 0.25 + 0.25
elif angle > 14.5: # thick
return (30 - angle) / 15.5 * 0.25 + 0.5
else: # very thick
return 1 - angle / 14.5 * 0.25
return None
p_cue_top, p_cue_bottom, cue_length = cue
cue_ball, object_ball = CO_balls
## calculate where to aim based on fractional aiming technique
p_aim = object_ball[0] # set initial aim point to object ball center
# because the angle should be angle between pocket-to-OB line and aim line, but not CTC (center to center) line,
# we need to iterate several times to get the correct fraction
for iteration in xrange(3): # 3 is an arbitrary number here
dist_pocket2aim_line = zc.calc_triangle_area(
cue_ball[0], p_aim, pocket) * 2 / zc.euc_dist(cue_ball[0], p_aim)
dist_pocket2aim_point = zc.euc_dist(p_aim, pocket)
angle = np.arcsin(dist_pocket2aim_line / dist_pocket2aim_point)
angle = angle / 1.8 # adjust angle because user is not looking straight down. This adjustment is very rough
angle = angle / np.pi * 180 # convert to degrees
fraction = _angle2fraction(angle)
if pocket[0] < object_ball[0][0]: # cut to left
p_aim = (object_ball[0][0] + (1 - fraction) * object_ball[1] * 2,
object_ball[0][1])
else: # cut to right
p_aim = (object_ball[0][0] - (1 - fraction) * object_ball[1] * 2,
object_ball[0][1])
## calculate where is aimed
p_aimed = ((object_ball[0][1] - cue_ball[0][1]) /
(cue_ball[0][1] - p_cue_bottom[1]) *
(cue_ball[0][0] - p_cue_bottom[0]) + cue_ball[0][0],
object_ball[0][1])
rtn_msg = {'status': 'success'}
return (rtn_msg, (p_aim, p_aimed))
def find_pingpong(img, img_prev, mask_table, mask_ball_prev, rotation_matrix):
def get_ball_stat(mask_ball):
cnt_ball = zc.mask2cnt(mask_ball)
area = cv2.contourArea(cnt_ball)
center = cnt_ball.mean(axis = 0)[0]
center_homo = np.hstack((center, 1)).reshape(3, 1)
center_rotated = np.dot(rotation_matrix, center_homo)
return (area, center_rotated)
rtn_msg = {'status' : 'success'}
mask_ball = zc.color_inrange(img, 'HSV', H_L = 165, H_U = 25, S_L = 60, V_L = 90, V_U = 240)
mask_ball, _ = zc.get_small_blobs(mask_ball, max_area = 2300)
mask_ball, _ = zc.get_big_blobs(mask_ball, min_area = 8)
mask_ball, counter = zc.get_square_blobs(mask_ball, th_diff = 0.2, th_area = 0.2)
if counter == 0:
rtn_msg = {'status' : 'fail', 'message' : "No good color candidate"}
return (rtn_msg, None)
cnt_table = zc.mask2cnt(mask_table)
loc_table_center = zc.get_contour_center(cnt_table)[::-1]
mask_ball_ontable = np.bitwise_and(mask_ball, mask_table)
mask_ball_ontable = zc.get_closest_blob(mask_ball_ontable, loc_table_center)
if mask_ball_ontable is not None: # if any ball on the table, we don't have to rely on previous ball positions
mask_ball = mask_ball_ontable
return (rtn_msg, (mask_ball, get_ball_stat(mask_ball)))
if mask_ball_prev is None: # mask_ball_ontable is already None
rtn_msg = {'status' : 'fail', 'message' : "Cannot initialize a location of ball"}
return (rtn_msg, None)
cnt_ball_prev = zc.mask2cnt(mask_ball_prev)
loc_ball_prev = zc.get_contour_center(cnt_ball_prev)[::-1]
mask_ball = zc.get_closest_blob(mask_ball, loc_ball_prev)
cnt_ball = zc.mask2cnt(mask_ball)
loc_ball = zc.get_contour_center(cnt_ball)[::-1]
ball_moved_dist = zc.euc_dist(loc_ball_prev, loc_ball)
if ball_moved_dist > 110:
rtn_msg = {'status' : 'fail', 'message' : "Lost track of ball: %d" % ball_moved_dist}
return (rtn_msg, None)
return (rtn_msg, (mask_ball, get_ball_stat(mask_ball)))
def find_table(img, display_list):
## find white border
DoB = zc.get_DoB(img, 1, 31, method='Average')
zc.check_and_display('DoB',
DoB,
display_list,
resize_max=config.DISPLAY_MAX_PIXEL,
wait_time=config.DISPLAY_WAIT_TIME)
mask_white = zc.color_inrange(DoB, 'HSV', V_L=10)
zc.check_and_display_mask('mask_white_raw',
img,
mask_white,
display_list,
resize_max=config.DISPLAY_MAX_PIXEL,
wait_time=config.DISPLAY_WAIT_TIME)
## find purple table (roughly)
#mask_table = zc.color_inrange(img, 'HSV', H_L = 130, H_U = 160, S_L = 50, V_L = 50, V_U = 220)
mask_ground = zc.color_inrange(img,
'HSV',
H_L=18,
H_U=30,
S_L=75,
S_U=150,
V_L=100,
V_U=255)
zc.check_and_display_mask('ground',
img,
mask_ground,
display_list,
resize_max=config.DISPLAY_MAX_PIXEL,
wait_time=config.DISPLAY_WAIT_TIME)
# red car
mask_red = zc.color_inrange(img, 'HSV', H_L=170, H_U=10, S_L=150)
mask_ground = np.bitwise_or(mask_ground, mask_red)
# ceiling
mask_ceiling = np.zeros((360, 640), dtype=np.uint8)
mask_ceiling[:40, :] = 255
mask_ground = np.bitwise_or(mask_ground, mask_ceiling)
# find the screen
mask_screen1 = zc.color_inrange(img,
'HSV',
H_L=15,
H_U=45,
S_L=30,
S_U=150,
V_L=40,
V_U=150)
mask_screen2 = ((img[:, :, 2] - 5) > img[:, :, 0]).astype(np.uint8) * 255
mask_screen = np.bitwise_or(mask_screen1, mask_screen2)
mask_screen = np.bitwise_and(np.bitwise_not(mask_ground), mask_screen)
mask_screen = zc.shrink(mask_screen, 5, iterations=3)
zc.check_and_display_mask('screen_raw',
img,
mask_screen,
display_list,
resize_max=config.DISPLAY_MAX_PIXEL,
wait_time=config.DISPLAY_WAIT_TIME)
bool_screen = zc.mask2bool([mask_screen])[0]
c_pixels = img[bool_screen].astype(np.int8)
d_pixels = c_pixels[:, 2] - c_pixels[:, 0]
rb_diff = np.median(d_pixels)
print rb_diff
if rb_diff > 20:
print "Case 1"
mask_table1 = zc.color_inrange(img,
'HSV',
H_L=0,
H_U=115,
S_U=45,
V_L=35,
V_U=120)
mask_table2 = zc.color_inrange(img,
'HSV',
H_L=72,
H_U=120,
S_L=20,
S_U=60,
V_L=35,
V_U=150)
mask_table = np.bitwise_or(mask_table1, mask_table2)
mask_screen = zc.color_inrange(img,
'HSV',
H_L=15,
H_U=45,
S_L=60,
S_U=150,
V_L=40,
V_U=150)
elif rb_diff > 15:
print "Case 2"
mask_table1 = zc.color_inrange(img,
'HSV',
H_L=0,
H_U=115,
S_U=45,
V_L=35,
V_U=120)
mask_table2 = zc.color_inrange(img,
'HSV',
H_L=72,
H_U=120,
S_L=20,
S_U=60,
V_L=35,
V_U=150)
mask_table = np.bitwise_or(mask_table1, mask_table2)
mask_screen = zc.color_inrange(img,
'HSV',
H_L=15,
H_U=45,
S_L=35,
S_U=150,
V_L=40,
V_U=150)
else:
print "Case 3"
mask_table1 = zc.color_inrange(img,
'HSV',
H_L=0,
H_U=115,
S_U=20,
V_L=35,
V_U=115)
mask_table2 = zc.color_inrange(img,
'HSV',
H_L=72,
H_U=120,
S_L=20,
S_U=60,
V_L=35,
V_U=150)
mask_table = np.bitwise_or(mask_table1, mask_table2)
mask_screen1 = zc.color_inrange(img,
'HSV',
H_L=15,
H_U=45,
S_L=30,
S_U=150,
V_L=40,
V_U=150)
mask_screen2 = (
(img[:, :, 2] - 1) > img[:, :, 0]).astype(np.uint8) * 255
mask_screen = np.bitwise_or(mask_screen1, mask_screen2)
mask_screen = np.bitwise_and(np.bitwise_not(mask_ground), mask_screen)
zc.check_and_display_mask('screen',
img,
mask_screen,
display_list,
resize_max=config.DISPLAY_MAX_PIXEL,
wait_time=config.DISPLAY_WAIT_TIME)
mask_table = np.bitwise_and(np.bitwise_not(mask_screen), mask_table)
mask_table = np.bitwise_and(np.bitwise_not(zc.shrink(mask_ground, 3)),
mask_table)
mask_table, _ = zc.get_big_blobs(mask_table, min_area=50)
mask_table = cv2.morphologyEx(mask_table,
cv2.MORPH_CLOSE,
zc.generate_kernel(7, 'circular'),
iterations=1)
#mask_table, _ = zc.find_largest_CC(mask_table)
zc.check_and_display_mask('table_purple_raw',
img,
mask_table,
display_list,
resize_max=config.DISPLAY_MAX_PIXEL,
wait_time=config.DISPLAY_WAIT_TIME)
if mask_table is None:
rtn_msg = {'status': 'fail', 'message': 'Cannot find table'}
return (rtn_msg, None)
#mask_table_convex, _ = zc.make_convex(mask_table.copy(), app_ratio = 0.005)
#mask_table = np.bitwise_or(mask_table, mask_table_convex)
mask_table_raw = mask_table.copy()
zc.check_and_display_mask('table_purple',
img,
mask_table,
display_list,
resize_max=config.DISPLAY_MAX_PIXEL,
wait_time=config.DISPLAY_WAIT_TIME)
mask_table_convex, _ = zc.make_convex(zc.shrink(mask_table,
5,
iterations=5),
app_ratio=0.01)
mask_table_shrunk = zc.shrink(mask_table_convex, 5, iterations=3)
zc.check_and_display_mask('table_shrunk',
img,
mask_table_shrunk,
display_list,
resize_max=config.DISPLAY_MAX_PIXEL,
wait_time=config.DISPLAY_WAIT_TIME)
## fine tune the purple table based on white border
mask_white = np.bitwise_and(np.bitwise_not(mask_table_shrunk), mask_white)
if 'mask_white' in display_list:
gray = np.float32(mask_white)
dst = cv2.cornerHarris(gray, 10, 3, 0.04)
dst = cv2.dilate(dst, None)
img_white = img.copy()
img_white[mask_white > 0, :] = [0, 255, 0]
img_white[dst > 2.4e7] = [0, 0, 255]
zc.check_and_display('mask_white',
img_white,
display_list,
resize_max=config.DISPLAY_MAX_PIXEL,
wait_time=config.DISPLAY_WAIT_TIME)
#mask_table, _ = zc.make_convex(mask_table, app_ratio = 0.005)
for i in xrange(15):
mask_table = zc.expand(mask_table, 3)
mask_table = np.bitwise_and(np.bitwise_not(mask_white), mask_table)
mask_table, _ = zc.find_largest_CC(mask_table)
if mask_table is None:
rtn_msg = {
'status': 'fail',
'message': 'Cannot find table, case 2'
}
return (rtn_msg, None)
if i % 4 == 3:
mask_table, _ = zc.make_convex(mask_table, app_ratio=0.01)
#img_display = img.copy()
#img_display[mask_table > 0, :] = [0, 0, 255]
#zc.display_image('table%d-b' % i, img_display, resize_max = config.DISPLAY_MAX_PIXEL, wait_time = config.DISPLAY_WAIT_TIME)
#mask_white = np.bitwise_and(np.bitwise_not(mask_table), mask_white)
mask_table = np.bitwise_and(np.bitwise_not(mask_white), mask_table)
mask_table, _ = zc.find_largest_CC(mask_table)
mask_table, hull_table = zc.make_convex(mask_table, app_ratio=0.01)
zc.check_and_display_mask('table_purple_fixed',
img,
mask_table,
display_list,
resize_max=config.DISPLAY_MAX_PIXEL,
wait_time=config.DISPLAY_WAIT_TIME)
## check if table is big enough
table_area = cv2.contourArea(hull_table)
table_area_percentage = float(table_area) / img.shape[0] / img.shape[1]
if table_area_percentage < 0.06:
rtn_msg = {
'status': 'fail',
'message': "Detected table too small: %f" % table_area_percentage
}
return (rtn_msg, None)
## find top line of table
hull_table = np.array(zc.sort_pts(hull_table[:, 0, :], order_first='y'))
ul = hull_table[0]
ur = hull_table[1]
if ul[0] > ur[0]:
t = ul
ul = ur
ur = t
i = 2
# the top two points in the hull are probably on the top line, but may not be the corners
while i < hull_table.shape[0] and hull_table[i, 1] - hull_table[0, 1] < 80:
pt_tmp = hull_table[i]
if pt_tmp[0] < ul[0] or pt_tmp[0] > ur[0]:
# computing the area of the part of triangle that lies inside the table
triangle = np.vstack([pt_tmp, ul, ur]).astype(np.int32)
mask_triangle = np.zeros_like(mask_table)
cv2.drawContours(mask_triangle, [triangle], 0, 255, -1)
pts = mask_table_raw[mask_triangle.astype(bool)]
if np.sum(pts == 255) > 10:
break
if pt_tmp[0] < ul[0]:
ul = pt_tmp
else:
ur = pt_tmp
i += 1
else:
break
ul = [int(x) for x in ul]
ur = [int(x) for x in ur]
if 'table' in display_list:
img_table = img.copy()
img_table[mask_table.astype(bool), :] = [255, 0, 255]
#cv2.line(img_table, tuple(ul), tuple(ur), [0, 255, 0], 3)
zc.check_and_display('table',
img_table,
display_list,
resize_max=config.DISPLAY_MAX_PIXEL,
wait_time=config.DISPLAY_WAIT_TIME)
## sanity checks about table top line detection
if zc.euc_dist(ul, ur)**2 * 3.1 < table_area:
rtn_msg = {
'status':
'fail',
'message':
"Table top line too short: %f, %f" %
(zc.euc_dist(ul, ur)**2 * 3.1, table_area)
}
return (rtn_msg, None)
if abs(zc.line_angle(ul, ur)) > 0.4:
rtn_msg = {
'status': 'fail',
'message': "Table top line tilted too much"
}
return (rtn_msg, None)
# check if two table sides form a reasonable angle
mask_table_bottom = mask_table.copy()
mask_table_bottom[:-30] = 0
p_left_most = zc.get_edge_point(mask_table_bottom, (-1, 0))
p_right_most = zc.get_edge_point(mask_table_bottom, (1, 0))
if p_left_most is None or p_right_most is None:
rtn_msg = {
'status': 'fail',
'message': "Table doesn't occupy bottom part of image"
}
return (rtn_msg, None)
left_side_angle = zc.line_angle(ul, p_left_most)
right_side_angle = zc.line_angle(ur, p_right_most)
angle_diff = zc.angle_dist(left_side_angle,
right_side_angle,
angle_range=math.pi * 2)
if abs(angle_diff) > 2.0:
rtn_msg = {
'status': 'fail',
'message': "Angle between two side edge not right: %f" % angle_diff
}
return (rtn_msg, None)
if 'table' in display_list:
img_table = img.copy()
img_table[mask_table.astype(bool), :] = [255, 0, 255]
cv2.line(img_table, tuple(ul), tuple(ur), [0, 255, 0], 3)
zc.check_and_display('table',
img_table,
display_list,
resize_max=config.DISPLAY_MAX_PIXEL,
wait_time=config.DISPLAY_WAIT_TIME)
## rotate to make opponent upright, use table edge as reference
pts1 = np.float32(
[ul, ur, [ul[0] + (ur[1] - ul[1]), ul[1] - (ur[0] - ul[0])]])
pts2 = np.float32([[0, config.O_IMG_HEIGHT],
[config.O_IMG_WIDTH, config.O_IMG_HEIGHT], [0, 0]])
M = cv2.getAffineTransform(pts1, pts2)
img[np.bitwise_not(zc.get_mask(img, rtn_type="bool", th=3)), :] = [3, 3, 3]
img_rotated = cv2.warpAffine(img, M,
(config.O_IMG_WIDTH, config.O_IMG_HEIGHT))
## sanity checks about rotated opponent image
bool_img_rotated_valid = zc.get_mask(img_rotated, rtn_type="bool")
if float(bool_img_rotated_valid.sum()
) / config.O_IMG_WIDTH / config.O_IMG_HEIGHT < 0.6:
rtn_msg = {
'status':
'fail',
'message':
"Valid area too small after rotation: %f" %
(float(bool_img_rotated_valid.sum()) / config.O_IMG_WIDTH /
config.O_IMG_HEIGHT)
}
return (rtn_msg, None)
rtn_msg = {'status': 'success'}
return (rtn_msg, (img_rotated, mask_table, M))
def find_pingpong(img, img_prev, mask_table, mask_ball_prev, rotation_matrix,
display_list):
def get_ball_stat(mask_ball):
cnt_ball = zc.mask2cnt(mask_ball)
area = cv2.contourArea(cnt_ball)
center = cnt_ball.mean(axis=0)[0]
center_homo = np.hstack((center, 1)).reshape(3, 1)
center_rotated = np.dot(rotation_matrix, center_homo)
return (area, center_rotated)
rtn_msg = {'status': 'success'}
mask_ball = zc.color_inrange(img,
'HSV',
H_L=165,
H_U=25,
S_L=65,
V_L=150,
V_U=255)
mask_screen = find_screen(img, display_list)
mask_screen, _ = zc.find_largest_CC(mask_screen)
mask_possible = np.bitwise_or(mask_screen, zc.expand(mask_table, 3))
mask_ball = np.bitwise_and(mask_ball, mask_possible)
mask_ball, _ = zc.get_small_blobs(mask_ball, max_area=2300)
mask_ball, _ = zc.get_big_blobs(mask_ball, min_area=8)
mask_ball, counter = zc.get_square_blobs(mask_ball,
th_diff=0.2,
th_area=0.2)
zc.check_and_display_mask('ball_raw',
img,
mask_ball,
display_list,
resize_max=config.DISPLAY_MAX_PIXEL,
wait_time=config.DISPLAY_WAIT_TIME)
if counter == 0:
rtn_msg = {'status': 'fail', 'message': "No good color candidate"}
return (rtn_msg, None)
cnt_table = zc.mask2cnt(mask_table)
loc_table_center = zc.get_contour_center(cnt_table)[::-1]
mask_ball_ontable = np.bitwise_and(mask_ball, mask_table)
mask_ball_ontable = zc.get_closest_blob(mask_ball_ontable,
loc_table_center)
if mask_ball_ontable is not None: # if any ball on the table, we don't have to rely on previous ball positions
mask_ball = mask_ball_ontable
zc.check_and_display_mask('ball',
img,
mask_ball,
display_list,
resize_max=config.DISPLAY_MAX_PIXEL,
wait_time=config.DISPLAY_WAIT_TIME)
return (rtn_msg, (mask_ball, get_ball_stat(mask_ball)))
if mask_ball_prev is None: # mask_ball_ontable is already None
rtn_msg = {
'status': 'fail',
'message': "Cannot initialize a location of ball"
}
return (rtn_msg, None)
cnt_ball_prev = zc.mask2cnt(mask_ball_prev)
loc_ball_prev = zc.get_contour_center(cnt_ball_prev)[::-1]
mask_ball = zc.get_closest_blob(mask_ball, loc_ball_prev)
zc.check_and_display_mask('ball',
img,
mask_ball,
display_list,
resize_max=config.DISPLAY_MAX_PIXEL,
wait_time=config.DISPLAY_WAIT_TIME)
cnt_ball = zc.mask2cnt(mask_ball)
loc_ball = zc.get_contour_center(cnt_ball)[::-1]
ball_moved_dist = zc.euc_dist(loc_ball_prev, loc_ball)
if ball_moved_dist > 110:
rtn_msg = {
'status': 'fail',
'message': "Lost track of ball: %d" % ball_moved_dist
}
return (rtn_msg, None)
return (rtn_msg, (mask_ball, get_ball_stat(mask_ball)))
def _detect_cue(img, mask_tables, mask_balls, display_list):
CUE_MIN_LENGTH = int(float(img.shape[1]) / 640 * 40 + 0.5)
PARA1 = int(float(img.shape[1]) / 640 * 2 + 0.5)
mask_blue, mask_bluer, mask_table, mask_table_fat = mask_tables
### edges on the table
#img_table = np.zeros(img.shape, dtype=np.uint8)
#img_table = cv2.bitwise_and(img, img, dst = img_table, mask = mask_table_convex)
#bw_table = cv2.cvtColor(img_table, cv2.COLOR_BGR2GRAY)
#edge_table = cv2.Canny(bw_table, 80, 160)
#edge_table = zc.expand(edge_table, 2)
#zc.check_and_display("edge_table", edge_table, display_list, resize_max = config.DISPLAY_MAX_PIXEL, wait_time = config.DISPLAY_WAIT_TIME)
### detect cue
#lines = cv2.HoughLinesP(edge_table, 1, np.pi/180, 30, minLineLength = 70, maxLineGap = 3)
#if lines is None:
# rtn_msg = {'status': 'fail', 'message' : 'Cannot find cue'}
# return (rtn_msg, None)
#lines = lines[0]
#if 'cue_edge' in display_list:
# img_cue = img.copy()
# for line in lines:
# pt1 = (line[0], line[1])
# pt2 = (line[2], line[3])
# cv2.line(img_cue, pt1, pt2, (255, 0, 255), 2)
# zc.check_and_display("cue_edge", img_cue, display_list, resize_max = config.DISPLAY_MAX_PIXEL, wait_time = config.DISPLAY_WAIT_TIME)
## interesting parts on the table (pockets, cue, hand, etc.)
mask_table_convex, _ = zc.make_convex(mask_table.copy(),
use_approxPolyDp=False)
zc.check_and_display_mask("table_convex",
img,
mask_table_convex,
display_list,
resize_max=config.DISPLAY_MAX_PIXEL,
wait_time=config.DISPLAY_WAIT_TIME)
mask_interesting = cv2.subtract(
cv2.subtract(mask_table_convex, mask_table), mask_bluer)
mask_interesting = cv2.subtract(mask_interesting, mask_balls)
mask_interesting = zc.shrink(mask_interesting, PARA1)
zc.check_and_display_mask("interesting",
img,
mask_interesting,
display_list,
resize_max=config.DISPLAY_MAX_PIXEL,
wait_time=config.DISPLAY_WAIT_TIME)
# find the blob with cue (and probably hand)
# TODO: this may be more robust with find_largest_CC function, in the case of half ball close to the bottom
mask_cue_hand = zc.get_closest_blob(
mask_interesting.copy(), (img.shape[0], img.shape[1] / 2),
min_length=CUE_MIN_LENGTH,
hierarchy_req='outer') # cue must be close to the bottom
## find cue top
p_cue_top = zc.get_edge_point(mask_cue_hand, (0, -1))
if p_cue_top is None:
rtn_msg = {'status': 'fail', 'message': 'Cannot find cue top'}
return (rtn_msg, None)
## find cue bottom
# the cue detected initially may not have reached the bottom of the image
for i in xrange(10):
mask_cue_hand = zc.expand_with_bound(mask_cue_hand,
cv2.bitwise_not(mask_bluer))
mask_cue_bottom = mask_cue_hand.copy()
mask_cue_bottom[:-2, :] = 0
mask_cue_bottom[:, :p_cue_top[0] - 40] = 0
mask_cue_bottom[:, p_cue_top[0] + 40:] = 0
nonzero = np.nonzero(mask_cue_bottom)
if len(nonzero) < 2 or len(nonzero[0]) == 0:
rtn_msg = {'status': 'fail', 'message': 'Cannot find cue bottom'}
return (rtn_msg, None)
rows, cols = nonzero
p_cue_bottom = ((np.min(cols) + np.max(cols)) / 2, img.shape[0])
## cue info
cue_length = zc.euc_dist(p_cue_top, p_cue_bottom)
if 'cue' in display_list:
img_cue = img.copy()
img_cue[mask_cue_hand > 0, :] = [0, 255, 255]
cv2.circle(img_cue, p_cue_top, 3, (255, 0, 255), -1)
cv2.line(img_cue, p_cue_top, p_cue_bottom, (255, 0, 255), 2)
zc.display_image("cue",
img_cue,
resize_max=config.DISPLAY_MAX_PIXEL,
wait_time=config.DISPLAY_WAIT_TIME)
## skeletonize
#skeleton_cue_hand = zc.skeletonize(mask_cue_hand)
rtn_msg = {'status': 'success'}
return (rtn_msg, (p_cue_top, p_cue_bottom, cue_length))
def find_table(img, o_img_height, o_img_width):
## find white border
DoB = zc.get_DoB(img, 1, 31, method = 'Average')
mask_white = zc.color_inrange(DoB, 'HSV', V_L = 10)
## find purple table (roughly)
mask_table = zc.color_inrange(img, 'HSV', H_L = 130, H_U = 160, S_L = 50, V_L = 50, V_U = 220)
mask_table, _ = zc.get_big_blobs(mask_table, min_area = 50)
mask_table = cv2.morphologyEx(mask_table, cv2.MORPH_CLOSE, zc.generate_kernel(7, 'circular'), iterations = 1)
mask_table, _ = zc.find_largest_CC(mask_table)
if mask_table is None:
rtn_msg = {'status': 'fail', 'message' : 'Cannot find table'}
return (rtn_msg, None)
mask_table_convex, _ = zc.make_convex(mask_table.copy(), app_ratio = 0.005)
mask_table = np.bitwise_or(mask_table, mask_table_convex)
mask_table_raw = mask_table.copy()
## fine tune the purple table based on white border
mask_white = np.bitwise_and(np.bitwise_not(mask_table), mask_white)
for i in range(15):
mask_table = zc.expand(mask_table, 3)
mask_table = np.bitwise_and(np.bitwise_not(mask_white), mask_table)
if i % 4 == 3:
mask_table, _ = zc.make_convex(mask_table, app_ratio = 0.01)
mask_table = np.bitwise_and(np.bitwise_not(mask_white), mask_table)
mask_table, _ = zc.find_largest_CC(mask_table)
mask_table, hull_table = zc.make_convex(mask_table, app_ratio = 0.01)
## check if table is big enough
table_area = cv2.contourArea(hull_table)
table_area_percentage = float(table_area) / img.shape[0] / img.shape[1]
if table_area_percentage < 0.06:
rtn_msg = {'status' : 'fail', 'message' : "Detected table too small: %f" % table_area_percentage}
return (rtn_msg, None)
## find top line of table
hull_table = np.array(zc.sort_pts(hull_table[:,0,:], order_first = 'y'))
ul = hull_table[0]
ur = hull_table[1]
if ul[0] > ur[0]:
t = ul; ul = ur; ur = t
i = 2
# the top two points in the hull are probably on the top line, but may not be the corners
while i < hull_table.shape[0] and hull_table[i, 1] - hull_table[0, 1] < 80:
pt_tmp = hull_table[i]
if pt_tmp[0] < ul[0] or pt_tmp[0] > ur[0]:
# computing the area of the part of triangle that lies inside the table
triangle = np.vstack([pt_tmp, ul, ur]).astype(np.int32)
mask_triangle = np.zeros_like(mask_table)
cv2.drawContours(mask_triangle, [triangle], 0, 255, -1)
pts = mask_table_raw[mask_triangle.astype(bool)]
if np.sum(pts == 255) > 10:
break
if pt_tmp[0] < ul[0]:
ul = pt_tmp
else:
ur = pt_tmp
i += 1
else:
break
ul = [int(x) for x in ul]
ur = [int(x) for x in ur]
## sanity checks about table top line detection
if zc.euc_dist(ul, ur) ** 2 * 2.5 < table_area:
rtn_msg = {'status' : 'fail', 'message' : "Table top line too short"}
return (rtn_msg, None)
if abs(zc.line_angle(ul, ur)) > 0.4:
rtn_msg = {'status' : 'fail', 'message' : "Table top line tilted too much"}
return (rtn_msg, None)
# check if two table sides form a reasonable angle
mask_table_bottom = mask_table.copy()
mask_table_bottom[:-30] = 0
p_left_most = zc.get_edge_point(mask_table_bottom, (-1, 0))
p_right_most = zc.get_edge_point(mask_table_bottom, (1, 0))
if p_left_most is None or p_right_most is None:
rtn_msg = {'status' : 'fail', 'message' : "Table doesn't occupy bottom part of image"}
return (rtn_msg, None)
left_side_angle = zc.line_angle(ul, p_left_most)
right_side_angle = zc.line_angle(ur, p_right_most)
angle_diff = zc.angle_dist(left_side_angle, right_side_angle, angle_range = math.pi * 2)
if abs(angle_diff) > 1.8:
rtn_msg = {'status' : 'fail', 'message' : "Angle between two side edge not right"}
return (rtn_msg, None)
## rotate to make opponent upright, use table edge as reference
pts1 = np.float32([ul,ur,[ul[0] + (ur[1] - ul[1]), ul[1] - (ur[0] - ul[0])]])
pts2 = np.float32([[0, o_img_height], [o_img_width, o_img_height], [0, 0]])
M = cv2.getAffineTransform(pts1, pts2)
img[np.bitwise_not(zc.get_mask(img, rtn_type = "bool", th = 3)), :] = [3,3,3]
img_rotated = cv2.warpAffine(img, M, (o_img_width, o_img_height))
## sanity checks about rotated opponent image
bool_img_rotated_valid = zc.get_mask(img_rotated, rtn_type = "bool")
if float(bool_img_rotated_valid.sum()) / o_img_width / o_img_height < 0.7:
rtn_msg = {'status' : 'fail', 'message' : "Valid area too small after rotation"}
return (rtn_msg, None)
rtn_msg = {'status' : 'success'}
return (rtn_msg, (img_rotated, mask_table, M))