[feat] Add shape operation tools (Layout-Parser#72)

* Add shape operation tools

* Add test for tools

* Add scipy as the dep

Author: lolipopshock

setup.py

@@ -37,6 +37,7 @@
       install_requires=[
         "numpy", 
         "opencv-python",
+        "scipy",
         "pandas",
         "pillow",
         "pyyaml>=5.1",

src/layoutparser/__init__.py

@@ -52,6 +52,11 @@
         "is_pytesseract_available",
         "is_gcv_available",
         "requires_backends"
+    ],
+    "tools": [
+        "generalized_connected_component_analysis_1d",
+        "simple_line_detection",
+        "group_textblocks_based_on_category"
     ]
 }
 

src/layoutparser/tools/__init__.py

@@ -0,0 +1,5 @@
+from .shape_operations import (
+    generalized_connected_component_analysis_1d,
+    simple_line_detection,
+    group_textblocks_based_on_category,
+)

src/layoutparser/tools/shape_operations.py

@@ -0,0 +1,167 @@
+# Copyright 2021 The Layout Parser team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List, Union, Any, Callable, Iterable
+from functools import partial, reduce
+
+import numpy as np
+from scipy.sparse import csr_matrix
+from scipy.sparse.csgraph import connected_components
+
+from ..elements import BaseLayoutElement, TextBlock
+
+
+def generalized_connected_component_analysis_1d(
+    sequence: List[Any],
+    scoring_func: Callable[[Any, Any], int],
+    aggregation_func: Callable[[List[Any]], Any] = None,
+    default_score_value: int = 0,
+) -> List[Any]:
+    """Perform connected componenet analysis for any 1D sequence based on
+    the scoring function and the aggregation function.
+    It will generate the adjacency_matrix for the 1D sequence object using
+    the provided `scoring_func` and find the connected componenets.
+    The `aggregation_func` will be used to aggregate all elements within
+    identified components (when not set, it will be the identity function).
+
+    Args:
+        sequence (List[Any]):
+            The provided 1D sequence of objects.
+        scoring_func (Callable[[Any, Any], int]):
+            The scoring function used to construct the adjacency_matrix.
+            It should take two objects in the sequence and produe a integer.
+        aggregation_func (Callable[[List[Any]], Any], optional):
+            The function used to aggregate the elements within an identified
+            component.
+            Defaults to the identify function: `lambda x: x`.
+        default_score_value (int, optional):
+            Used to set the default (background) score values that should be
+            not considered when running connected component analysis.
+            Defaults to 0.
+
+    Returns:
+        List[Any]: A list of length n - the number of the detected componenets.
+    """
+
+    if aggregation_func is None:
+        aggregation_func = lambda x: x  # Identity Function
+
+    seq_len = len(sequence)
+    adjacency_matrix = np.ones((seq_len, seq_len)) * default_score_value
+
+    for i in range(seq_len):
+        for j in range(i + 1, seq_len):
+            adjacency_matrix[i][j] = scoring_func(sequence[i], sequence[j])
+
+    graph = csr_matrix(adjacency_matrix)
+    n_components, labels = connected_components(
+        csgraph=graph, directed=False, return_labels=True
+    )
+
+    grouped_sequence = []
+    for comp_idx in range(n_components):
+        element_idx = np.where(labels == comp_idx)[0]
+        grouped_sequence.append(aggregation_func([sequence[i] for i in element_idx]))
+
+    return grouped_sequence
+
+
+def simple_line_detection(
+    layout: Iterable[BaseLayoutElement], x_tolerance: int = 10, y_tolerance: int = 10
+) -> List[BaseLayoutElement]:
+    """Perform line detection based on connected component analysis.
+
+    The is_line_wise_close is the scoring function, which returns True
+    if the y-difference is smaller than the y_tolerance AND the
+    x-difference (the horizontal gap between two boxes) is also smaller
+    than the x_tolerance, and False otherwise.
+
+    All the detected components will then be passed into aggregation_func,
+    which returns the overall union box of all the elements, or the line
+    box.
+
+    Args:
+        layout (Iterable):
+            A list (or Layout) of BaseLayoutElement
+        x_tolerance (int, optional):
+            The value used for specifying the maximum allowed y-difference
+            when considered whether two tokens are from the same line.
+            Defaults to 10.
+        y_tolerance (int, optional):
+            The value used for specifying the maximum allowed horizontal gap
+            when considered whether two tokens are from the same line.
+            Defaults to 10.
+
+    Returns:
+        List[BaseLayoutElement]: A list of BaseLayoutElement, denoting the line boxes.
+    """
+
+    def is_line_wise_close(token_a, token_b, x_tolerance, y_tolerance):
+        y_a = token_a.block.center[1]
+        y_b = token_b.block.center[1]
+
+        a_left, a_right = token_a.block.coordinates[0::2]
+        b_left, b_right = token_b.block.coordinates[0::2]
+
+        return (
+            abs(y_a - y_b) <= y_tolerance
+            and min(abs(a_left - b_right), abs(a_right - b_left)) <= x_tolerance
+        )
+        # If the y-difference is smaller than the y_tolerance AND
+        # the x-difference (the horizontal gap between two boxes)
+        # is also smaller than the x_tolerance threshold, then
+        # these two tokens are considered as line-wise close.
+
+    detected_lines = generalized_connected_component_analysis_1d(
+        layout,
+        scoring_func=partial(
+            is_line_wise_close, y_tolerance=x_tolerance, x_tolerance=y_tolerance
+        ),
+        aggregation_func=lambda seq: reduce(layout[0].__class__.union, seq),
+    )
+
+    return detected_lines
+
+
+def group_textblocks_based_on_category(
+    layout: Iterable[TextBlock], union_group: bool = True
+) -> Union[List[TextBlock], List[List[TextBlock]]]:
+    """Group textblocks based on their category (block.type).
+
+    Args:
+        layout (Iterable):
+            A list (or Layout) of BaseLayoutElement
+        union_group (bool):
+            Whether to union the boxes within each group.
+            Defaults to True.
+
+    Returns:
+        List[TextBlock]: When `union_group=True`, it produces a list of
+            TextBlocks, denoting the boundaries of each texblock group.
+        List[List[TextBlock]]: When `union_group=False`, it preserves
+            the elements within each group for further processing.
+    """
+
+    if union_group:
+        aggregation_func = lambda seq: reduce(layout[0].__class__.union, seq)
+    else:
+        aggregation_func = None
+
+    detected_group_boxes = generalized_connected_component_analysis_1d(
+        layout,
+        scoring_func=lambda a, b: a.type == b.type,
+        aggregation_func=aggregation_func,
+    )
+
+    return detected_group_boxes

tests/test_tools.py

@@ -0,0 +1,68 @@
+# Copyright 2021 The Layout Parser team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from layoutparser import load_pdf
+from layoutparser.tools import (
+    generalized_connected_component_analysis_1d,
+    simple_line_detection,
+    group_textblocks_based_on_category,
+)
+
+def test_generalized_connected_component_analysis_1d():
+    
+    A = [1, 2, 3]
+    
+    results = generalized_connected_component_analysis_1d(
+        A,
+        scoring_func=lambda x,y: abs(x-y)<=1
+    )
+    assert len(results) == 1
+    
+    A = [1, 2, 3, 5, 6, 7]
+    results = generalized_connected_component_analysis_1d(
+        A,
+        scoring_func=lambda x,y: abs(x-y)<=1
+    )
+    assert len(results) == 2
+    
+    A = [1, 2, 3, 5, 6, 7]
+    results = generalized_connected_component_analysis_1d(
+        A,
+        scoring_func=lambda x,y: abs(x-y)<=2
+    )
+    assert len(results) == 1
+    
+    A = [1, 2, 3, 5, 6, 7]
+    results = generalized_connected_component_analysis_1d(
+        A,
+        scoring_func=lambda x,y: abs(x-y)<=1,
+        aggregation_func=max
+    )
+    assert results == [3, 7]
+    
+def test_simple_line_detection():
+    
+    page_layout = load_pdf("tests/fixtures/io/example.pdf")[0]
+    
+    pdf_lines = simple_line_detection(page_layout)
+    
+    assert len(pdf_lines) == 15
+    
+def test_group_textblocks_based_on_category():
+    
+    page_layout = load_pdf("tests/fixtures/io/example.pdf")[0]
+    
+    pdf_blocks = group_textblocks_based_on_category(page_layout)
+    
+    assert len(pdf_blocks) == 3