crop.py

class NotDiamondDashException(Exception):
	pass

def crop_dd_screenshot(pixarray):
	"""
	If there is no stored index, load the reference .png,
	find it in the current screen, and get the coordinates.

	If the reference is not found, raise a
	NotDiamondDashException so we know we're not looking
	at a Diamond Dash screen

	If the index has already been stored, just slice the
	array to the right parameters.
	"""
	global TOP_LEFT_INDEX

	if not TOP_LEFT_INDEX:
        # Load it up
		ref_pixarray = read_png_to_pixarray('topleft_ref.png')

		ind = search_for_subarray(pixarray, ref_pixarray)

		if not ind:
			raise NotDiamondDashException("This isn't diamond dash.")


		ref_row, ref_col, _ = ref_pixarray.shape

		# Add the size of the reference
		row = ind[0] + ref_row
		col = ind[1] + ref_col

		TOP_LEFT_INDEX = (row, col)
	else:
		row, col = TOP_LEFT_INDEX

	# Diamond dash board runs 400x360 px.
	return pixarray[row:row + 360, col:col + 400], (row, col)

def read_png_to_pixarray(filename):
	"""
	Read a .png image into an (x,y,3) full-color pixarray
	"""
	with open(filename, 'rb') as f:
		r = png.Reader(file=f)
		out = r.asDirect()

		pixarray = numpy.reshape(list(out[2]), (out[1], out[0], 3))

	return pixarray

def search_for_subarray(A, A_sub):
	"""
	Search for A_sub in A, and return the coordinates to it

	>>> A = numpy.array([[1,2,3,4],[5,6,7,8]])
	>>> search_for_subarray(A, numpy.array([[2,3],[6,7]]))
	(0, 1)
	"""

	if len(A.shape) == 3:
		sub_rows, sub_cols, _ = A_sub.shape
		rows, cols, _ = A.shape
	elif len(A.shape) == 2:
		sub_rows, sub_cols = A_sub.shape
		rows, cols = A.shape

	for i in range(rows):
		for j in range(cols):
			if numpy.all(A[i][j] == A_sub[0][0]):
				if numpy.all(A[i:i+sub_rows, j:j+sub_cols] == A_sub):
					return (i, j)
	return None