fix: max instead of mean structure size for line kernels

Author: OBrink

decimer_segmentation/complete_structure.py

@@ -368,7 +368,7 @@ def get_neighbour_pixels(
 
 def detect_horizontal_and_vertical_lines(
     image: np.ndarray,
-    average_depiction_size: Tuple[int, int]
+    max_depiction_size: Tuple[int, int]
 ) -> np.ndarray:
     """
     This function takes an image and returns a binary mask that labels the pixels that
@@ -378,15 +378,16 @@ def detect_horizontal_and_vertical_lines(
     Args:
         image (np.ndarray): binarised image (np.array; type bool) as it is returned by
             binary_erosion() in complete_structure_mask()
+        max_depiction_size (Tuple[int, int]): height, width; used as thresholds
 
     Returns:
         np.ndarray: Exclusion mask that contains indices of pixels that are part of
             horizontal or vertical lines
     """
     binarised_im = ~image * 255
     binarised_im = binarised_im.astype("uint8")
-    
-    structure_height, structure_width = average_depiction_size
+
+    structure_height, structure_width = max_depiction_size
 
     horizontal_kernel = cv2.getStructuringElement(
         cv2.MORPH_RECT, (structure_width, 1)
@@ -477,23 +478,23 @@ def expansion_coordination(
 def complete_structure_mask(
     image_array: np.array,
     mask_array: np.array,
-    average_depiction_size: Tuple[int, int],
+    max_depiction_size: Tuple[int, int],
     debug=False
 ) -> np.array:
     """
     This funtion takes an image (np.array) and an array containing the masks (shape:
     x,y,n where n is the amount of masks and x and y are the pixel coordinates).
-    Additionally, it takes the average depiction size of the structures in the image
+    Additionally, it takes the maximal depiction size of the structures in the image
     which is used to define the kernel size for the vertical and horizontal line
     detection for the exclusion masks. The exclusion mask is used to exclude pixels
-    from the mask expansion to avoid including whole tables. 
+    from the mask expansion to avoid including whole tables.
     It detects objects on the contours of the mask and expands it until it frames the
     complete object in the image. It returns the expanded mask array
 
     Args:
         image_array (np.array): input image
         mask_array (np.array): shape: y, x, n where n is the amount of masks
-        average_depiction_size (Tuple[int, int]): height, width
+        max_depiction_size (Tuple[int, int]): height, width
         debug (bool, optional): More verbose if True. Defaults to False.
 
     Returns:
@@ -519,7 +520,7 @@ def complete_structure_mask(
             [mask_array[:, :, index] for index in range(mask_array.shape[2])]
         )
         exclusion_mask = detect_horizontal_and_vertical_lines(blurred_image_array,
-                                                              average_depiction_size)
+                                                              max_depiction_size)
         # Run expansion the expansion
         image_repeat = itertools.repeat(blurred_image_array, mask_array.shape[2])
         exclusion_mask_repeat = itertools.repeat(exclusion_mask, mask_array.shape[2])

decimer_segmentation/decimer_segmentation.py

@@ -117,11 +117,11 @@ def segment_chemical_structures(
     return segments
 
 
-def determine_average_depiction_size(
+def determine_depiction_size_with_buffer(
     bboxes: List[Tuple[int, int, int, int]]
 ) -> Tuple[int, int]:
     """
-    This function takes a list of bounding boxes and returns the average
+    This function takes a list of bounding boxes and returns 1.2 * the maximal
     depiction size (height, width) of the depicted chemical structures.
 
     Args:
@@ -138,9 +138,9 @@ def determine_average_depiction_size(
         width = bbox[3] - bbox[1]
         heights.append(height)
         widths.append(width)
-    average_height = int(np.mean(heights))
-    average_width = int(np.mean(widths))
-    return int(average_height), int(average_width)
+    height = int(1.2 * np.max(heights))
+    width = int(1.2 * np.max(widths))
+    return height, width
 
 
 def sort_segments_bboxes(
@@ -227,11 +227,11 @@ def get_expanded_masks(image: np.array) -> np.array:
     """
     # Structure detection with MRCNN
     masks, bboxes, _ = get_mrcnn_results(image)
-    size = determine_average_depiction_size(bboxes)
+    size = determine_depiction_size_with_buffer(bboxes)
     # Mask expansion
     expanded_masks = complete_structure_mask(image_array=image,
                                              mask_array=masks,
-                                             average_depiction_size=size,)
+                                             max_depiction_size=size,)
     return expanded_masks
 
 

tests/test_decimer_segmentation.py

@@ -1,13 +1,13 @@
 import numpy as np
 from decimer_segmentation.decimer_segmentation import *
 
-def test_determine_average_depiction_size():
+def test_determine_depiction_size_with_buffer():
     # Determine the average depiction size of the structures in a given list of structures
     bboxes = [
         [0, 0, 6, 6],
         [0, 0, 8, 8],
         [0, 0, 10, 10],
     ]
-    expected_result = (8, 8)
-    actual_result = determine_average_depiction_size(bboxes)
+    expected_result = (12, 12)
+    actual_result = determine_depiction_size_with_buffer(bboxes)
     assert expected_result == actual_result

tests/test_mask_expansion.py

@@ -153,7 +153,7 @@ def test_detect_horizontal_and_vertical_lines():
         [[False] * 20] * 20
     )
     expected_result[9] = np.array([True] * 20)
-    actual_result = detect_horizontal_and_vertical_lines(~test_image)
+    actual_result = detect_horizontal_and_vertical_lines(~test_image, (2, 2))
     np.testing.assert_array_equal(expected_result, actual_result)