def find_cueballs(image, table):
"""Finds N best cueballs, evaluates them an returns a list of the
cueballs sorted by best to worst."""
global CURRENT_FRAME
matchmap = matchTemplate(image, tmpl_ball)
cueballs = []
for i in range(NUM_BALLS):
minval, maxval, minloc, maxloc = cvMinMaxLoc(matchmap)
if DEBUG: print("Found %dth cueball at min %d:%d = %.3f, max %d:%d = %.3f" %
(i, minloc.x + int(tmpl_ball.width/2), minloc.y, minval, maxloc.x, maxloc.y, maxval))
# Ball position is set as the top center of the ball template
ball = Cueball(minloc.x + int(tmpl_ball.width/2), minloc.y, minval)
ball.flags['nth_match'] = i
ball.confidence = evaluate_cueball(ball, image, table)
cueballs.append(ball)
# Remove the local minimum so we can use cvMinMaxLoc again
removeMinimum(matchmap, minloc.x, minloc.y, tmpl_ball.width, tmpl_ball.height)
# We are able to save a pretty interesting normalized heightmap of the
# match here
#cvSaveImage(here("results/%d_hmap_ball.jpg" % CURRENT_FRAME), normalize(matchmap))
cueballs = sorted(cueballs, key=lambda x: x.confidence)
cueballs.reverse()
return cueballs
def removeMinimum(img, x, y, w, h):
"""Utility function for removing the global minimum from a cvArray.
This finds the global maximum of the same array and then blits a rectangle of the same color
over the area designated by x,y and w,h.
This way to find the next _local_ minimum, we can just remove the current global
minimum and search for the global one again using cvMinMaxLoc - should be faster than looking
for all local minima with python."""
minmax = cvMinMaxLoc(img)
cvRectangle(img, cvPoint(x-int(w/2), y-int(h/2)), cvPoint(x+int(w/2), y+int(h/2)), cvScalar(minmax[1], 0, 0, 0), CV_FILLED)
def get_table(image):
"""Find the table in the image by matching each pocket separately.
Uses POCKET_TRESHOLD and some sanity checks to filter matches.
Returns a dictionary with keys 'tl', 'tr', 'bl', 'br', with a cvPoint instance
as the value for each one.
"""
bounds = {}
for pocket in ['tl', 'tr', 'bl', 'br']:
matchmap = matchTemplate(image, pocket_templates[pocket])
minval, maxval, minloc, maxloc = cvMinMaxLoc(matchmap)
#print("Pocket %s value = %d" % (pocket, minval))
if minval < POCKET_TRESHOLD:
# Found a possible match
bounds[pocket] = cvPoint(minloc.x + pocket_markers[pocket].x, minloc.y + pocket_markers[pocket].y)
else:
# If the best match for a pocket is above the treshold, then we can discard
# the whole thing
return None
if DEBUG: print("Possible table bounds %s" % bounds)
# Sanity check
if bounds['tl'].x >= bounds['tr'].x or bounds['bl'].x >= bounds['br'].x:
return None
if bounds['tl'].x <= bounds['bl'].x or bounds['tr'].x >= bounds['br'].x:
return None
if bounds['bl'].y <= bounds['tl'].y:
return None
# Make sure top and bottom edges are straight
if abs( bounds['tl'].y - bounds['tr'].y ) > image.height / 80:
return None
if abs( bounds['bl'].y - bounds['br'].y ) > image.height / 80:
return None
# Make sure table is more or less symmetrical
if math.fabs((bounds['tl'].x-bounds['bl'].x)-(bounds['br'].x-bounds['tr'].x)) > image.width / 100:
return None
return bounds