import cv2 import numpy as np from threshold import preprocess from utils import find_corners, draw_circle_at_corners, grid_line_helper, draw_line from utils import clean_square_helper, classify_one_digit #----------------Process pipe line------------------------------# # 1) Threshold Adaptive to get gray-scale image to find contours # 2) Find contours from original image # 3) Image alignment (warp image) on original image # 4) Get horizontal, vertical line and create grid mask # 5) Extract numbers and split gray-scale image into 81 squares # 6) Clean noise pixels of each square # 7) Recognize digits # 8) Solve sudoku # 9) Draw solved board on warped image # 10) Unwarped image --> Result def find_contours(img, original): """ contours: A tuple of all point creating contour lines, each contour is a np array of points (x,y). hierachy: [Next, Previous, First_Child, Parent] contour approximation: https://pyimagesearch.com/2021/10/06/opencv-contour-approximation/ """ # find contours on threshold image contours, hierachy = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) #sort the largest contour to find the puzzle contours = sorted(contours, key = cv2.contourArea, reverse = True) polygon = None # find the largest rectangle-shape contour to make sure this is the puzzle for con in contours: area = cv2.contourArea(con) perimeter = cv2.arcLength(con, closed = True) approx = cv2.approxPolyDP(con, epsilon=0.01 * perimeter, closed = True) num_of_ptr = len(approx) if num_of_ptr == 4 and area > 1000: polygon = con #finded puzzle break if polygon is not None: # find corner top_left = find_corners(polygon, limit_func= min, compare_func= np.add) top_right = find_corners(polygon, limit_func= max, compare_func= np.subtract) bot_left = find_corners(polygon,limit_func=min, compare_func= np.subtract) bot_right = find_corners(polygon,limit_func=max, compare_func=np.add) #Check polygon is square, if not return [] #Set threshold rate for width and height to determine square bounding box if not (0.5 < ((top_right[0]-top_left[0]) / (bot_right[1]-top_right[1]))<1.5): print("Exception 1 : Get another image to get square-shape puzzle") return [],[],[] if bot_right[1] - top_right[1] == 0: print("Exception 2 : Get another image to get square-shape puzzle") return [],[],[] corner_list = [top_left, top_right, bot_right, bot_left] draw_original = original.copy() cv2.drawContours(draw_original, [polygon], 0, (0,255,0), 3) #draw circle at each corner point for x in corner_list: draw_circle_at_corners(draw_original, x) return draw_original, corner_list, original # draw_original: Img which drown contour and corner # corner_list: list of 4 corner points # original: Original imgs print("Can not detect puzzle") return [],[],[] def warp_image(corner_list, original): """ Input: 4 corner points and threshold grayscale image Output: Perspective transformation matrix and transformed image Perspective transformation: https://theailearner.com/tag/cv2-warpperspective/ """ try: corners = np.array(corner_list, dtype= "float32") top_left, top_right, bot_left, bot_right = corners[0], corners[1], corners[2], corners[3] #Get the largest side to be the side of squared transfromed puzzle side = int(max([ np.linalg.norm(top_right - bot_right), np.linalg.norm(top_left - bot_left), np.linalg.norm(bot_right - bot_left), np.linalg.norm(top_left - top_right) ])) out_ptr = np.array([[0,0],[side-1,0],[side-1,side-1], [0,side-1]],dtype="float32") transfrom_matrix = cv2.getPerspectiveTransform(corners, out_ptr) transformed_image = cv2.warpPerspective(original, transfrom_matrix, (side, side)) return transformed_image, transfrom_matrix except IndexError: print("Can not detect corners") except: print("Something went wrong. Try another image") def get_grid_line(img, length = 10): """ Get horizontal and vertical lines from warped image """ horizontal = grid_line_helper(img, shape_location= 1) vertical = grid_line_helper(img, shape_location=0) return vertical, horizontal def create_grid_mask(horizontal, vertical): """ Completely detect all lines by using Hough Transformation Create grid mask to extract number by using bitwise_and with warped images """ # combine two line to make a grid grid = cv2.add(horizontal, vertical) # Apply threshold to cover more area # grid = cv2.adaptiveThreshold(grid, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 235, 2) morpho_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3,3)) grid = cv2.dilate(grid, morpho_kernel, iterations=2) # find the line by Houghline transfromation lines = cv2.HoughLines(grid, 0.3, np.pi/90, 200) lines_img = draw_line(grid, lines) # Extract all the lines mask = cv2.bitwise_not(lines_img) return mask def split_squares(number_img): """ Split number img into 81 squares. """ square_list = [] side = number_img.shape[0] // 9 #find each square and append to square_list for j in range(0,9): for i in range(0,9): top_left_square = (i * side, j * side) bot_right_square = ((i+1) * side, (j+1) * side) square_list.append(number_img[top_left_square[1]:bot_right_square[1], top_left_square[0]: bot_right_square[0]]) return square_list def clean_square(square_list): """ Return cleaned-square list and number of digits available in the image Clean-square list has both 0 and images """ cleaned_squares = [] count = 0 for sq in square_list: new_img, is_num = clean_square_helper(sq) if is_num: cleaned_squares.append(new_img) count += 1 else: cleaned_squares.append(0) return cleaned_squares, count def clean_square_all_images(square_list): """ Return cleaned-square list Clean-square list has all images(images with no number with be black image after cleaning) """ square_cleaned_list = [] for i in square_list: clean_square, _ = clean_square_helper(i) square_cleaned_list.append(clean_square) return square_cleaned_list def recognize_digits(model, resized, org_img): res_str = "" for img in resized: digit = classify_one_digit(model, img, org_img) res_str += str(digit) return res_str def draw_digits_on_warped(warped_img, solved_board, unsolved_board): """ Function to draw digits from solved board to warped img """ width = warped_img.shape[0] // 9 img_w_text = np.zeros_like(warped_img) for j in range(9): for i in range(9): if unsolved_board[j][i] == 0: # Only draw new number to blank cell in warped image, avoid overlapping p1 = (i * width, j * width) # Top left corner of a bounding box p2 = ((i + 1) * width, (j + 1) * width) # Bottom right corner of bounding box # Find the center of square to draw digit center = ((p1[0] + p2[0]) // 2, (p1[1] + p2[1]) // 2) text_size, _ = cv2.getTextSize(str(solved_board[j][i]), cv2.FONT_HERSHEY_SIMPLEX, 1, 6) text_origin = (center[0] - text_size[0] // 2, center[1] + text_size[1] // 2) cv2.putText(warped_img, str(solved_board[j][i]), text_origin, cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 0, 255), 6) return img_w_text, warped_img def unwarp_image(img_src, img_dest, pts, time): pts = np.array(pts) height, width = img_src.shape[0], img_src.shape[1] pts_source = np.array([[0, 0], [width - 1, 0], [width - 1, height - 1], [0, width - 1]], dtype='float32') matrix, status = cv2.findHomography(pts_source, pts) # Covert to original view perspective warped = cv2.warpPerspective(img_src, matrix, (img_dest.shape[1], img_dest.shape[0])) # Draw a black rectangle in img_dest cv2.fillConvexPoly(img_dest, pts, 0, 16) dst_img = cv2.add(img_dest, warped) dst_img_height, dst_img_width = dst_img.shape[0], dst_img.shape[1] cv2.putText(dst_img, "Time solved: {} s".format(str(np.round(time,4))), (dst_img_width - 360, 60), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2) return dst_img