-
Notifications
You must be signed in to change notification settings - Fork 0
/
utils.py
192 lines (139 loc) · 5.14 KB
/
utils.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
import cv2
import numpy as np
import operator
import torch
def find_corners(polygon, limit_func, compare_func):
"""
Input: Rectangle puzzle extract from contours
Output: One of four cornet point depend on limit_func, compare_func
# limit_fn is the min or max function
# compare_fn is the np.add or np.subtract function
Note: (0,0) point is at the top-left
top-left: (x+y) min
top-right: (x-y) max
bot-left: (x-y) min
bot-right: (x+y) max
"""
index, _ = limit_func(enumerate([compare_func(ptr[0][0], ptr[0][1]) for ptr in polygon]), key = operator.itemgetter(1))
return polygon[index][0][0], polygon[index][0][1]
def draw_circle_at_corners(original, ptr):
"""
Helper function to draw circle at corners
"""
cv2.circle(original, ptr, 5, (0,255,0), cv2.FILLED)
def grid_line_helper(img, shape_location, length = 10):
"""
Helper function to fine vertical, horizontal line
Find horizontal line: shape_location = 0
Find vertical line: shape_location = 1
"""
clone_img = img.copy()
row_or_col = clone_img.shape[shape_location]
# Find the distance between lines
size = row_or_col // length
# Morphological transfromation to find line
# Define morphology kernel
if shape_location == 0:
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (size,1))
else:
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1,size))
clone_img = cv2.erode(clone_img, kernel)
clone_img = cv2.dilate(clone_img, kernel)
return clone_img
def draw_line(img, lines):
"""
Draw all lines in lines got from cv2.HoughLine()
"""
clone_img = img.copy()
# lines list from cv2.HoughLine() is 3d array
# Convert to 2d array
lines = np.squeeze(lines)
for rho, theta in lines:
a = np.cos(theta)
b = np.sin(theta)
x0 = a*rho
y0 = b*rho
x1 = int(x0 + 1000 * (-b))
y1 = int(y0 + 1000 * a)
x2 = int(x0 - 1000 * (-b))
y2 = int(y0 - 1000 * a)
#Draw line every loop
cv2.line(clone_img, (x1,y1), (x2,y2), (255,255,255), 4)
return clone_img
def clean_square_helper(img):
"""
Clean noises in every square splited
Input: One of 81 squares
Output: Cleaned square and boolean var which so that there is number in it
"""
if np.isclose(img, 0).sum() / (img.shape[0] * img.shape[1]) >= 0.96:
return np.zeros_like(img), False
# if there is very little white in the region around the center, this means we got an edge accidently
height, width = img.shape
mid = width // 2
if np.isclose(img[:, int(mid - width * 0.38):int(mid + width * 0.38)], 0).sum() / (2 * width * 0.38 * height) >= 0.98:
return np.zeros_like(img), False
# center image
contours, _ = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
contours = sorted(contours, key=cv2.contourArea, reverse=True)
x, y, w, h = cv2.boundingRect(contours[0])
start_x = (width - w) // 2
start_y = (height - h) // 2
new_img = np.zeros_like(img)
new_img[start_y:start_y + h, start_x:start_x + w] = img[y:y + h, x:x + w]
return new_img, True
def resize_square(clean_square_list):
"""
Resize clean squares into 28x28 in order to feed to classifier
"""
resized_list = []
for img in clean_square_list:
resized = cv2.resize(img, (28,28), interpolation=cv2.INTER_AREA)
resized_list.append(resized)
return resized_list
def resize_square32(clean_square_list):
"""
Resize clean squares into 32x32 in order to feed to tf classifier
"""
resized_list = []
for img in clean_square_list:
resized = cv2.resize(img, (32,32), interpolation=cv2.INTER_AREA)
resized_list.append(resized)
return resized_list
def classify_one_digit(model, resize_square, org_image):
"""
Determine whether each square has number by counting number of (not black) pixel and compare to threshold value
Using classifier to predict number in square
- Return 0 if the square is blank
- Return predict digit if the square has number
"""
threshold = 0
if (org_image.shape[0] > 600 or org_image.shape[1] > 600) or (org_image.shape[1] > 600 or org_image.shape[2] > 600):
threshold = 40
else:
threshold = 60
# Determine blank square
if (resize_square != resize_square.min()).sum() < threshold:
return str(0)
model.eval()
# Convert to shape (1,1,28,28) to be compatible with dataloader for evaluation
iin = torch.Tensor(resize_square).unsqueeze(0).unsqueeze(0)
with torch.no_grad():
out = model(iin)
# Get index of predict digit
_, index = torch.max(out, 1)
pred_digit = index.item()
return str(pred_digit)
def normalize(resized_list):
"""
Scale pixel value for recognition
"""
return [img/255 for img in resized_list]
def convert_str_to_board(string, step = 9):
"""
Convert recognized string into 2D array for sudoku solving
"""
board = []
for i in range(0, len(string), step):
board.append([int(char) for char in string[i:i+step]])
return np.array(board)