Refactoring and simplification.

home/xword.py

@@ -14,6 +14,26 @@ def preprocess_image(original, gaussian_blur_size, adaptive_threshold_block_size
     return img
 
 
+def morph_open_image(img, kernel_size, iterations=1):
+    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, kernel_size)
+    return cv2.morphologyEx(img, cv2.MORPH_OPEN, kernel, iterations=iterations)
+
+
+def get_fundamental_frequency(fft):
+    mag = abs(fft[0:len(fft) // 2])
+    mag[0] = 0
+    return int(np.argmax(mag))
+
+
+def get_line_fft(img, line_detector_element_size, axis):
+    lines = morph_open_image(img, (line_detector_element_size, 1) if axis == 1 else (1, line_detector_element_size))
+    return np.fft.fft(np.sum(lines, axis=axis))
+
+
+def get_line_frequency(img, line_detector_element_size, axis):
+    return get_fundamental_frequency(get_line_fft(img, line_detector_element_size, axis))
+
+
 def find_biggest_contour(img):
     contours, hierarchy = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
 
@@ -28,13 +48,10 @@ def find_biggest_contour(img):
     return biggest
 
 
-def erode_contour(img_shape, contour, kernel_size, iterations):
+def erode_contour(img_shape, contour, erosion_kernel_size, iterations):
     contour_img = np.zeros(img_shape, dtype=np.uint8)
     cv2.drawContours(contour_img, [contour], 0, 255, -1)
-
-    kernel = np.ones((kernel_size, kernel_size), dtype=np.uint8)
-    contour_img = cv2.erode(contour_img, kernel, iterations=iterations)
-    contour_img = cv2.dilate(contour_img, kernel, iterations=iterations)
+    contour_img = morph_open_image(contour_img, (erosion_kernel_size, erosion_kernel_size), iterations)
     return find_biggest_contour(contour_img)
 
 
@@ -87,12 +104,6 @@ def extract_square(img, top_left, top_right, bottom_right, bottom_left):
     return cv2.warpPerspective(img, m, (int(longest), int(longest)))
 
 
-def get_fundamental_frequency(fft):
-    mag = abs(fft[0:len(fft) // 2])
-    mag[0] = 0
-    return int(np.argmax(mag))
-
-
 def get_threshold_from_quantile(img, quantile):
     height, width = img.shape
     num_pixels = height * width
@@ -174,7 +185,7 @@ def extract_crossword(
     square=True,
     num_dilations=1,
     contour_erosion_kernel_size=5,
-    contour_erosion_iterations=5,
+    contour_erosion_iterations=6,
     line_detector_element_size=51,
     sampling_block_size_ratio=0.25,
     sampling_threshold_quantile=0.3,
@@ -198,19 +209,8 @@ def extract_crossword(
 
     img = extract_square(img, top_left, top_right, bottom_right, bottom_left)
 
-    horiz_elem = cv2.getStructuringElement(cv2.MORPH_RECT, (line_detector_element_size, 1))
-    horiz_lines = cv2.erode(img, horiz_elem)
-    horiz_lines = cv2.dilate(horiz_lines, horiz_elem)
-
-    vert_elem = cv2.getStructuringElement(cv2.MORPH_RECT, (1, line_detector_element_size))
-    vert_lines = cv2.erode(img, vert_elem)
-    vert_lines = cv2.dilate(vert_lines, vert_elem)
-
-    row_fft = np.fft.fft(np.sum(horiz_lines, axis=1))
-    col_fft = np.fft.fft(np.sum(vert_lines, axis=0))
-
-    num_rows = get_fundamental_frequency(row_fft)
-    num_cols = get_fundamental_frequency(col_fft)
+    num_rows = get_line_frequency(img, line_detector_element_size, 1)
+    num_cols = get_line_frequency(img, line_detector_element_size, 0)
     if square and (num_rows != num_cols):
         warnings.append("Crossword is not square")