custom-layoutparser 0.1.0__py3-none-any.whl

layoutparser/ocr/tesseract_agent.py

@@ -0,0 +1,193 @@
+# 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.
+
+import io
+import csv
+import pickle
+
+import pandas as pd
+
+from .base import BaseOCRAgent, BaseOCRElementType
+from ..io import load_dataframe
+from ..file_utils import is_pytesseract_available
+
+if is_pytesseract_available():
+    import pytesseract
+
+
+class TesseractFeatureType(BaseOCRElementType):
+    """
+    The element types for Tesseract Detection API
+    """
+
+    PAGE = 0
+    BLOCK = 1
+    PARA = 2
+    LINE = 3
+    WORD = 4
+
+    @property
+    def attr_name(self):
+        name_cvt = {
+            TesseractFeatureType.PAGE: "page_num",
+            TesseractFeatureType.BLOCK: "block_num",
+            TesseractFeatureType.PARA: "par_num",
+            TesseractFeatureType.LINE: "line_num",
+            TesseractFeatureType.WORD: "word_num",
+        }
+        return name_cvt[self]
+
+    @property
+    def group_levels(self):
+        levels = ["page_num", "block_num", "par_num", "line_num", "word_num"]
+        return levels[: self + 1]
+
+
+class TesseractAgent(BaseOCRAgent):
+    """
+    A wrapper for `Tesseract <https://github.com/tesseract-ocr/tesseract>`_ Text
+    Detection APIs based on `PyTesseract <https://github.com/tesseract-ocr/tesseract>`_.
+    """
+
+    DEPENDENCIES = ["pytesseract"]
+
+    def __init__(self, languages="eng", **kwargs):
+        """Create a Tesseract OCR Agent.
+
+        Args:
+            languages (:obj:`list` or :obj:`str`, optional):
+                You can specify the language code(s) of the documents to detect to improve
+                accuracy. The supported language and their code can be found on
+                `its github repo <https://github.com/tesseract-ocr/langdata>`_.
+                It supports two formats: 1) you can pass in the languages code as a string
+                of format like `"eng+fra"`, or 2) you can pack them as a list of strings
+                `["eng", "fra"]`.
+                Defaults to 'eng'.
+        """
+        self.lang = languages if isinstance(languages, str) else "+".join(languages)
+        self.configs = kwargs
+
+    @classmethod
+    def with_tesseract_executable(cls, tesseract_cmd_path, **kwargs):
+
+        pytesseract.pytesseract.tesseract_cmd = tesseract_cmd_path
+        return cls(**kwargs)
+
+    def _detect(self, img_content):
+        res = {}
+        res["text"] = pytesseract.image_to_string(
+            img_content, lang=self.lang, **self.configs
+        )
+        _data = pytesseract.image_to_data(img_content, lang=self.lang, **self.configs)
+        res["data"] = pd.read_csv(
+            io.StringIO(_data),
+            quoting=csv.QUOTE_NONE,
+            encoding="utf-8",
+            sep="\t",
+            converters={"text": str},
+        )
+        return res
+
+    def detect(
+        self, image, return_response=False, return_only_text=True, agg_output_level=None
+    ):
+        """Send the input image for OCR.
+
+        Args:
+            image (:obj:`np.ndarray` or :obj:`str`):
+                The input image array or the name of the image file
+            return_response (:obj:`bool`, optional):
+                Whether directly return all output (string and boxes
+                info) from Tesseract.
+                Defaults to `False`.
+            return_only_text (:obj:`bool`, optional):
+                Whether return only the texts in the OCR results.
+                Defaults to `False`.
+            agg_output_level (:obj:`~TesseractFeatureType`, optional):
+                When set, aggregate the GCV output with respect to the
+                specified aggregation level. Defaults to `None`.
+        """
+
+        res = self._detect(image)
+
+        if return_response:
+            return res
+
+        if return_only_text:
+            return res["text"]
+
+        if agg_output_level is not None:
+            return self.gather_data(res, agg_output_level)
+
+        return res["text"]
+
+    @staticmethod
+    def gather_data(response, agg_level):
+        """
+        Gather the OCR'ed text, bounding boxes, and confidence
+        in a given aggeragation level.
+        """
+        assert isinstance(
+            agg_level, TesseractFeatureType
+        ), f"Invalid agg_level {agg_level}"
+        res = response["data"]
+        df = (
+            res[~res.text.isna()]
+            .groupby(agg_level.group_levels)
+            .apply(
+                lambda gp: pd.Series(
+                    [
+                        gp["left"].min(),
+                        gp["top"].min(),
+                        gp["width"].max(),
+                        gp["height"].max(),
+                        gp["conf"].mean(),
+                        gp["text"].str.cat(sep=" "),
+                    ]
+                )
+            )
+            .reset_index(drop=True)
+            .reset_index()
+            .rename(
+                columns={
+                    0: "x_1",
+                    1: "y_1",
+                    2: "w",
+                    3: "h",
+                    4: "score",
+                    5: "text",
+                    "index": "id",
+                }
+            )
+            .assign(
+                x_2=lambda x: x.x_1 + x.w,
+                y_2=lambda x: x.y_1 + x.h,
+                block_type="rectangle",
+            )
+            .drop(columns=["w", "h"])
+        )
+
+        return load_dataframe(df)
+
+    @staticmethod
+    def load_response(filename):
+        with open(filename, "rb") as fp:
+            res = pickle.load(fp)
+        return res
+
+    @staticmethod
+    def save_response(res, file_name):
+
+        with open(file_name, "wb") as fp:
+            pickle.dump(res, fp, protocol=pickle.HIGHEST_PROTOCOL)

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,
+)

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