-
Notifications
You must be signed in to change notification settings - Fork 0
/
processing.py
241 lines (186 loc) · 8.56 KB
/
processing.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
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