rubiks_scanner_core.py

import skimage.transform
import skimage.filters
import numpy as np
import cv2
from keras import backend as K
from keras.models import load_model
from skimage import img_as_ubyte


# Match an input image to a template image using SIFT and return the input image warped to the template image coordinates
def get_scorecard_sift(image, template):

	MIN_MATCH_COUNT = 10

	# Initiate SIFT detector
	sift = cv2.xfeatures2d.SIFT_create()

	# find the keypoints and descriptors with SIFT
	kp1, des1 = sift.detectAndCompute(image, None)
	kp2, des2 = sift.detectAndCompute(template, None)

	FLANN_INDEX_KDTREE = 0
	index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
	search_params = dict(checks=50)

	flann = cv2.FlannBasedMatcher(index_params, search_params)

	matches = flann.knnMatch(des1, des2, k=2)

	# store all the good matches as per Lowe's ratio test.
	good = []
	for m, n in matches:
		if m.distance < 0.55 * n.distance:
			good.append(m)

	print("# Matches: ", len(good))

	if len(good) > MIN_MATCH_COUNT:
		src_pts = np.float32([kp1[m.queryIdx].pt for m in good]).reshape(-1, 1, 2)
		dst_pts = np.float32([kp2[m.trainIdx].pt for m in good]).reshape(-1, 1, 2)

		M, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)
		matchesMask = mask.ravel().tolist()
		print(M)
		h, w = image.shape
		pts = np.float32([[0, 0], [0, h - 1], [w - 1, h - 1], [w - 1, 0]]).reshape(-1, 1, 2)
		dst = cv2.perspectiveTransform(pts, M)
		img2 = cv2.polylines(template, [np.int32(dst)], True, 255, 3, cv2.LINE_AA)

		h_2, w_2 = img2.shape
		adjusted_image = cv2.warpPerspective(image, M, (w_2, h_2))
		# cv2.imshow("Warped", adjusted_image)
		# cv2.waitKey(0)
		# cv2.destroyAllWindows()

		return adjusted_image
	else:
		print("Not enough matches are found - %d/%d" % (len(good), MIN_MATCH_COUNT))
		matchesMask = None
		return None


# Predict an array of images of digits
def predict_digits(digit_images, digit_flags):
	img_rows = 28
	img_cols = 28
	for i in range(len(digit_images)):
		if K.image_data_format() == 'channels_first':
			digit_images[i] = digit_images[i].reshape(1, img_rows, img_cols)
		else:
			digit_images[i] = digit_images[i].reshape(img_rows, img_cols, 1)
	model = load_model('CNN\\new_model.h5')
	print("# Digits:", len(digit_images))
	predictions = model.predict(np.asarray(digit_images))
	predicted_digits = []
	flags = []
	i = 0
	for prediction in predictions:
		which_digit = np.argmax(prediction)
		confidence = np.max(prediction)
		# print("Predicted", which_digit, "with confidence", confidence)
		predicted_digits.append(which_digit)
		if confidence < 0.75 or digit_flags[i] == 1:
			flags.append(1)
		else:
			flags.append(0)
		i += 1

	return predicted_digits, flags


# Extracts the ID from the scorecard
def get_id_from_scorecard(image):
	bw = image < skimage.filters.threshold_local(image, 101)
	bw = bw.astype("float32")
	digits = []
	for column in range(0, 3):
		min_y = 134
		max_y = 177
		min_x = 43 + 54 * column
		max_x = 90 + 54 * column
		digit = bw[min_y:max_y, min_x:max_x]
		digits.append(digit)
		# cv2.imshow("digit", digit)
		# cv2.waitKey(0)
		# cv2.destroyAllWindows()
	return digits


# Extract a digit from a candidate digit image by finding a contour that looks like a digit and finding its bounding
# box
def extract_digit(digit):
	cnts = cv2.findContours(img_as_ubyte(digit.copy()), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
	cnts = cnts[1]
	cnts = sorted(cnts, key=cv2.contourArea, reverse=True)[:2]
	added = False
	# Loop over all of the contour candidates
	for c in cnts:
		# Compute the bounding box of each contour
		(x, y, w, h) = cv2.boundingRect(c)
		# If the contour is sufficiently large, it should be a digit
		# A small lower bound on the width is due to the fact that 1s are small
		if not ((40 >= h >= 12) and (40 >= w >= 3)):
			continue
		added = True
		digit_crop = digit[y:y + h, x:x + w]
		resize_ratio = min(20. / w, 20. / h)
		resize_width = int(resize_ratio * w)
		resize_height = int(resize_ratio * h)
		# Resize the digit image to 20x20 according to its aspect ratio
		digit_resized = skimage.transform.resize(digit_crop, (resize_height, resize_width), mode='constant')
		digit_28_28 = np.zeros((28, 28), dtype=float)
		lower_bound_y = int((28 - resize_height) / 2)
		upper_bound_y = int(resize_height + (28 - resize_height) / 2)
		lower_bound_x = int((28 - resize_width) / 2)
		upper_bound_x = int(resize_width + (28 - resize_width) / 2)
		# Pad the 20x20 with 0s to make a 28x28 digit image with the digit centered in the image
		digit_28_28[lower_bound_y:upper_bound_y, lower_bound_x:upper_bound_x] = digit_resized
		return digit_28_28, 0
	if not added:
		# If a digit was not found, add on a blank image and a flag to indicate it isn't a digit
		return np.zeros((28, 28)), 1


# Extract each of the 35 digits from the scorecard using hardcoded coordinates from the template image we made
def get_digits_from_scorecard(image):
	bw = image < skimage.filters.threshold_local(image, 101)

	digits = []
	flags = []

	for column in range(0, 3):
		min_y = 134
		max_y = 177
		min_x = 43 + 54 * column
		max_x = 90 + 54 * column
		digit, flag = extract_digit(bw[min_y:max_y, min_x:max_x])
		digits.append(digit)
		flags.append(flag)

	for row in range(0, 5):
		for column in range(0, 7):
			min_y = 233 + 49 * row
			max_y = 273 + 49 * row
			min_x = 43 + 54 * column
			max_x = 90 + 54 * column
			# cv2.imshow("bw", np.float32(bw[min_y:max_y, min_x:max_x]))
			# cv2.waitKey(0)
			# cv2.destroyAllWindows()
			digit, flag = extract_digit(bw[min_y:max_y, min_x:max_x])
			# cv2.imshow("digit", np.float32(digit))
			# cv2.waitKey(0)
			# cv2.destroyAllWindows()
			digits.append(digit)
			flags.append(flag)

	return digits, flags


# Construct the ID from a list of three digits
def construct_id(digits):
	return str(digits[0]) + str(digits[1]) + str(digits[2])


# Construct each of the round times from a list of digits
def construct_times(digits):
	times = []
	for i in range(0, 5):
		round_time = str(digits[0 + 7*i]) + str(digits[1 + 7*i]) + ":" + str(digits[2 + 7*i]) + str(digits[3 + 7*i]) + ":" + str(digits[4 + 7*i]) + str(
			digits[5 + 7*i]) + str(digits[6 + 7*i])
		times.append(round_time)
	return times


# Find rectangles in the input image. If their aspect ratio is within a certain range, return true and draw the image
# on the input image. This is meant to give the use an idea of whether the scorecard is being picked up in the image.
def found_contour_of_template(image):
	gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
	gray = cv2.GaussianBlur(gray, (5, 5), 0)
	edged = cv2.Canny(gray, 20, 150)  # 75, 200
	# Find the contours in the edged image, keeping only the largest ones
	(_, cnts, _) = cv2.findContours(edged.copy(), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
	cnts = sorted(cnts, key=cv2.contourArea, reverse=True)[:5]

	rectangle = []
	# loop over the contours
	for c in cnts:
		# approximate the contour
		peri = cv2.arcLength(c, True)
		approx = cv2.approxPolyDP(c, 0.03 * peri, True)  # 0.02
		# if our approximated contour has four points, then we
		# can assume that we have found our screen
		if len(approx) == 4:
			rectangle = approx
			break
	if len(rectangle) > 0:
		point_1 = rectangle[0][0]
		point_2 = rectangle[1][0]
		point_3 = rectangle[2][0]

		dist_1_2 = np.sqrt((point_2[0] - point_1[0])**2 + (point_2[1] - point_1[1])**2)
		dist_2_3 = np.sqrt((point_3[0] - point_2[0])**2 + (point_3[1] - point_2[1])**2)

		ratio = 0
		if dist_1_2 > dist_2_3:
			ratio = dist_1_2 / dist_2_3
			print("Found ratio of", ratio)
		else:
			ratio = dist_2_3 / dist_1_2
			print("Found ratio of", ratio)
		if 1.1 <= ratio <= 1.6:
			cv2.drawContours(image, [rectangle], -1, (0, 255, 0), 2)
			return True
		return False
	return False