[WIP] Refactoring meta-detection

cv_analysis/layout_parsing.py

@@ -1,14 +1,15 @@
-from functools import reduce
+from functools import reduce, partial
 from itertools import compress
 from operator import __and__
 from typing import Iterable
 
 import cv2
 import numpy as np
-from funcy import lmap, compose
+from funcy import lmap, compose, rcompose
 
 from cv_analysis.utils.connect_rects import connect_related_rects2
 from cv_analysis.utils.conversion import box_to_rectangle, rectangle_to_box
+from cv_analysis.utils.display import show_image
 from cv_analysis.utils.postprocessing import (
     remove_included,
     has_no_parent,
@@ -20,28 +21,39 @@ def parse_layout(image: np.array):
 
     rectangles = find_segments(image)
     rectangles = remove_included(rectangles)
-    rectangles = lmap(rectangle_to_box, rectangles)
-    rectangles = connect_related_rects2(rectangles)
-    rectangles = lmap(box_to_rectangle, rectangles)
+    boxes = lmap(rectangle_to_box, rectangles)
+    boxes = connect_related_rects2(boxes)
+    rectangles = lmap(box_to_rectangle, boxes)
     rectangles = remove_included(rectangles)
 
     return rectangles
 
 
-def find_segments(image, meta=1):
+def find_segments(image):
+    rectangles = rcompose(
+        prepare_for_initial_detection,
+        __find_segments,
+        partial(prepare_for_meta_detection, image.copy()),
+        __find_segments,
+    )(image)
 
-    preprocess = compose(dilate_page_components, normalize_to_gray_scale)
+    return rectangles
 
-    contours, hierarchies = find_contours(preprocess(image) if meta else image)
+
+def prepare_for_initial_detection(image: np.ndarray):
+    return compose(dilate_page_components, normalize_to_gray_scale)(image)
+
+
+def __find_segments(image):
+
+    contours, hierarchies = find_contours(image)
     likely_segments = map(is_likely_segment, contours)
     without_parents = map(has_no_parent, hierarchies[0])
     valid_contours = map(__and__, likely_segments, without_parents)
     contours = compress(contours, valid_contours)
-    rectangles = lmap(compose(box_to_rectangle, cv2.boundingRect), contours)
 
-    if meta:
-        preprocessed = meta_detection(image.copy(), rectangles)
-        rectangles = find_segments(preprocessed, 0)
+    boxes = map(cv2.boundingRect, contours)
+    rectangles = lmap(box_to_rectangle, boxes)
 
     return rectangles
 
@@ -62,10 +74,7 @@ def dilate_page_components(image):
     return cv2.dilate(thresh, kernel, iterations=4)
 
 
-def meta_detection(image: np.ndarray, rectangles: Iterable[Rectangle]):
-    """Given a list of previously detected segments, rerun the detection algorithm. Heuristically this improves the
-    quality of the detection.
-    """
+def prepare_for_meta_detection(image: np.ndarray, rectangles: Iterable[Rectangle]):
 
     image = fill_rectangles(image, rectangles)
     image = threshold_image(image)