Spaces:
Running
Running
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 | |