Refactoring

- add typehints
- other minor refactorings

cv_analysis/figure_detection/text.py

@@ -1,5 +1,7 @@
 import cv2
 
+from cv_analysis.layout_parsing import normalize_to_gray_scale
+
 
 def remove_primary_text_regions(image):
     """Removes regions of primary text, meaning no figure descriptions for example, but main text body paragraphs.
@@ -35,6 +37,7 @@ def remove_primary_text_regions(image):
 
 def apply_threshold_to_image(image):
     """Converts an image to black and white."""
+    image = normalize_to_gray_scale(image)
     image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if len(image.shape) > 2 else image
     return cv2.threshold(image, 253, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1]
 

cv_analysis/layout_parsing.py

@@ -1,28 +1,29 @@
 from functools import reduce, partial
-from typing import Iterable
+from typing import Iterable, List
 
 import cv2
 import numpy as np
 from funcy import compose, rcompose, lkeep
 
 from cv_analysis.utils.common import find_contours
-from cv_analysis.utils.conversion import box_to_rectangle
+from cv_analysis.utils.conversion import box_to_rectangle, contour_to_rectangle
 from cv_analysis.utils.merging import connect_related_rectangles
 from cv_analysis.utils.postprocessing import remove_included, has_no_parent
 from cv_analysis.utils.rectangle import Rectangle
 
 
-def parse_layout(image: np.array):
-
-    rectangles = find_segments(image)
-    rectangles = remove_included(rectangles)
-    rectangles = connect_related_rectangles(rectangles)
-    rectangles = remove_included(rectangles)
+def parse_layout(image: np.array) -> List[Rectangle]:
+    rectangles = rcompose(
+        find_segments,
+        remove_included,
+        connect_related_rectangles,
+        remove_included,
+    )(image)
 
     return rectangles
 
 
-def find_segments(image):
+def find_segments(image: np.ndarray) -> List[Rectangle]:
     rectangles = rcompose(
         prepare_for_initial_detection,
         __find_segments,
@@ -33,29 +34,25 @@ def find_segments(image):
     return rectangles
 
 
-def prepare_for_initial_detection(image: np.ndarray):
+def prepare_for_initial_detection(image: np.ndarray) -> np.ndarray:
     return compose(dilate_page_components, normalize_to_gray_scale)(image)
 
 
-def __find_segments(image):
+def __find_segments(image: np.ndarray) -> List[Rectangle]:
     def to_rectangle_if_valid(contour, hierarchy):
-        return (
-            box_to_rectangle(cv2.boundingRect(contour))
-            if is_likely_segment(contour) and has_no_parent(hierarchy)
-            else None
-        )
+        return contour_to_rectangle(contour) if is_likely_segment(contour) and has_no_parent(hierarchy) else None
 
     rectangles = lkeep(map(to_rectangle_if_valid, *find_contours(image)))
 
     return rectangles
 
 
-def is_likely_segment(rect, min_area=100):
+def is_likely_segment(rectangle: Rectangle, min_area: float = 100) -> bool:
     # FIXME: Parameterize via factory
-    return cv2.contourArea(rect, False) > min_area
+    return cv2.contourArea(rectangle, False) > min_area
 
 
-def dilate_page_components(image):
+def dilate_page_components(image: np.ndarray) -> np.ndarray:
     # FIXME: Parameterize via factory
     image = cv2.GaussianBlur(image, (7, 7), 0)
     # FIXME: Parameterize via factory
@@ -63,10 +60,11 @@ def dilate_page_components(image):
     # FIXME: Parameterize via factory
     kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
     # FIXME: Parameterize via factory
-    return cv2.dilate(thresh, kernel, iterations=4)
+    dilate = cv2.dilate(thresh, kernel, iterations=4)
+    return dilate
 
 
-def prepare_for_meta_detection(image: np.ndarray, rectangles: Iterable[Rectangle]):
+def prepare_for_meta_detection(image: np.ndarray, rectangles: Iterable[Rectangle]) -> np.ndarray:
 
     image = fill_rectangles(image, rectangles)
     image = threshold_image(image)
@@ -76,28 +74,27 @@ def prepare_for_meta_detection(image: np.ndarray, rectangles: Iterable[Rectangle
     return image
 
 
-def normalize_to_gray_scale(image):
-    if len(image.shape) > 2:
-        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+def normalize_to_gray_scale(image: np.ndarray) -> np.ndarray:
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if len(image.shape) > 2 else image
     return image
 
 
-def threshold_image(image):
+def threshold_image(image: np.ndarray) -> np.ndarray:
     # FIXME: Parameterize via factory
     _, image = cv2.threshold(image, 254, 255, cv2.THRESH_BINARY)
     return image
 
 
-def invert_image(image):
+def invert_image(image: np.ndarray):
     return ~image
 
 
-def fill_rectangles(image, rectangles):
+def fill_rectangles(image: np.ndarray, rectangles: Iterable[Rectangle]) -> np.ndarray:
     image = reduce(fill_in_component_area, rectangles, image)
     return image
 
 
-def fill_in_component_area(image, rect):
+def fill_in_component_area(image: np.ndarray, rect: Rectangle) -> np.ndarray:
     x, y, w, h = rect
     cv2.rectangle(image, (x, y), (x + w, y + h), (0, 0, 0), -1)
     cv2.rectangle(image, (x, y), (x + w, y + h), (255, 255, 255), 7)

cv_analysis/utils/conversion.py

@@ -1,4 +1,5 @@
 import json
+from typing import Sequence
 
 import cv2
 
@@ -9,12 +10,12 @@ def contour_to_rectangle(contour):
     return box_to_rectangle(cv2.boundingRect(contour))
 
 
-def box_to_rectangle(box):
+def box_to_rectangle(box: Sequence[int]) -> Rectangle:
     x, y, w, h = box
     return Rectangle(x, y, x + w, y + h)
 
 
-def rectangle_to_box(rectangle):
+def rectangle_to_box(rectangle: Rectangle) -> Sequence[int]:
     return [rectangle.x1, rectangle.y1, rectangle.width, rectangle.height]
 
 

cv_analysis/utils/postprocessing.py

@@ -1,6 +1,6 @@
 from functools import partial
 from itertools import starmap, compress
-from typing import Iterable, List
+from typing import Iterable, List, Sequence
 
 from cv_analysis.utils.rectangle import Rectangle
 
@@ -17,8 +17,8 @@ def remove_overlapping(rectangles: Iterable[Rectangle]) -> List[Rectangle]:
 
 
 def remove_included(rectangles: Iterable[Rectangle]) -> List[Rectangle]:
-    keep = [rect for rect in rectangles if not rect.is_included(rectangles)]
-    return keep
+    rectangles_to_keep = [rect for rect in rectangles if not rect.is_included(rectangles)]
+    return rectangles_to_keep
 
 
 def __remove_isolated_unsorted(rectangles: Iterable[Rectangle]) -> List[Rectangle]:
@@ -45,9 +45,9 @@ def __remove_isolated_sorted(rectangles: Iterable[Rectangle]) -> List[Rectangle]
     return rectangles
 
 
-def remove_isolated(rectangles: Iterable[Rectangle], input_unsorted=True) -> List[Rectangle]:
+def remove_isolated(rectangles: Iterable[Rectangle], input_unsorted: bool = True) -> List[Rectangle]:
     return (__remove_isolated_unsorted if input_unsorted else __remove_isolated_sorted)(rectangles)
 
 
-def has_no_parent(hierarchy):
+def has_no_parent(hierarchy: Sequence[int]) -> bool:
     return hierarchy[-1] <= 0