Browse Source

Refactoring and simplification.

Andrew Klopper 6 years ago
parent
commit
82ef7f2999
1 changed files with 25 additions and 25 deletions
  1. 25 25
      home/xword.py

+ 25 - 25
home/xword.py

14
     return img
14
     return img
15
 
15
 
16
 
16
 
17
+def morph_open_image(img, kernel_size, iterations=1):
18
+    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, kernel_size)
19
+    return cv2.morphologyEx(img, cv2.MORPH_OPEN, kernel, iterations=iterations)
20
+
21
+
22
+def get_fundamental_frequency(fft):
23
+    mag = abs(fft[0:len(fft) // 2])
24
+    mag[0] = 0
25
+    return int(np.argmax(mag))
26
+
27
+
28
+def get_line_fft(img, line_detector_element_size, axis):
29
+    lines = morph_open_image(img, (line_detector_element_size, 1) if axis == 1 else (1, line_detector_element_size))
30
+    return np.fft.fft(np.sum(lines, axis=axis))
31
+
32
+
33
+def get_line_frequency(img, line_detector_element_size, axis):
34
+    return get_fundamental_frequency(get_line_fft(img, line_detector_element_size, axis))
35
+
36
+
17
 def find_biggest_contour(img):
37
 def find_biggest_contour(img):
18
     contours, hierarchy = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
38
     contours, hierarchy = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
19
 
39
 
28
     return biggest
48
     return biggest
29
 
49
 
30
 
50
 
31
-def erode_contour(img_shape, contour, kernel_size, iterations):
51
+def erode_contour(img_shape, contour, erosion_kernel_size, iterations):
32
     contour_img = np.zeros(img_shape, dtype=np.uint8)
52
     contour_img = np.zeros(img_shape, dtype=np.uint8)
33
     cv2.drawContours(contour_img, [contour], 0, 255, -1)
53
     cv2.drawContours(contour_img, [contour], 0, 255, -1)
34
-
35
-    kernel = np.ones((kernel_size, kernel_size), dtype=np.uint8)
36
-    contour_img = cv2.erode(contour_img, kernel, iterations=iterations)
37
-    contour_img = cv2.dilate(contour_img, kernel, iterations=iterations)
54
+    contour_img = morph_open_image(contour_img, (erosion_kernel_size, erosion_kernel_size), iterations)
38
     return find_biggest_contour(contour_img)
55
     return find_biggest_contour(contour_img)
39
 
56
 
40
 
57
 
87
     return cv2.warpPerspective(img, m, (int(longest), int(longest)))
104
     return cv2.warpPerspective(img, m, (int(longest), int(longest)))
88
 
105
 
89
 
106
 
90
-def get_fundamental_frequency(fft):
91
-    mag = abs(fft[0:len(fft) // 2])
92
-    mag[0] = 0
93
-    return int(np.argmax(mag))
94
-
95
-
96
 def get_threshold_from_quantile(img, quantile):
107
 def get_threshold_from_quantile(img, quantile):
97
     height, width = img.shape
108
     height, width = img.shape
98
     num_pixels = height * width
109
     num_pixels = height * width
174
     square=True,
185
     square=True,
175
     num_dilations=1,
186
     num_dilations=1,
176
     contour_erosion_kernel_size=5,
187
     contour_erosion_kernel_size=5,
177
-    contour_erosion_iterations=5,
188
+    contour_erosion_iterations=6,
178
     line_detector_element_size=51,
189
     line_detector_element_size=51,
179
     sampling_block_size_ratio=0.25,
190
     sampling_block_size_ratio=0.25,
180
     sampling_threshold_quantile=0.3,
191
     sampling_threshold_quantile=0.3,
198
 
209
 
199
     img = extract_square(img, top_left, top_right, bottom_right, bottom_left)
210
     img = extract_square(img, top_left, top_right, bottom_right, bottom_left)
200
 
211
 
201
-    horiz_elem = cv2.getStructuringElement(cv2.MORPH_RECT, (line_detector_element_size, 1))
202
-    horiz_lines = cv2.erode(img, horiz_elem)
203
-    horiz_lines = cv2.dilate(horiz_lines, horiz_elem)
204
-
205
-    vert_elem = cv2.getStructuringElement(cv2.MORPH_RECT, (1, line_detector_element_size))
206
-    vert_lines = cv2.erode(img, vert_elem)
207
-    vert_lines = cv2.dilate(vert_lines, vert_elem)
208
-
209
-    row_fft = np.fft.fft(np.sum(horiz_lines, axis=1))
210
-    col_fft = np.fft.fft(np.sum(vert_lines, axis=0))
211
-
212
-    num_rows = get_fundamental_frequency(row_fft)
213
-    num_cols = get_fundamental_frequency(col_fft)
212
+    num_rows = get_line_frequency(img, line_detector_element_size, 1)
213
+    num_cols = get_line_frequency(img, line_detector_element_size, 0)
214
     if square and (num_rows != num_cols):
214
     if square and (num_rows != num_cols):
215
         warnings.append("Crossword is not square")
215
         warnings.append("Crossword is not square")
216
 
216