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 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))