custom-layoutparser 0.1.0__py3-none-any.whl

layoutparser/models/paddledetection/layoutmodel.py

@@ -0,0 +1,297 @@
+# Copyright 2021 The Layout Parser team and Paddle Detection model 
+# contributors. 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, Dict, Any, Tuple
+import os
+from functools import reduce
+import warnings
+
+from PIL import Image
+import cv2
+import numpy as np
+
+from .catalog import PathManager, LABEL_MAP_CATALOG, MODEL_CATALOG
+from ..base_layoutmodel import BaseLayoutModel
+from ...elements import Rectangle, TextBlock, Layout
+
+from ...file_utils import is_paddle_available
+
+if is_paddle_available():
+    import paddle.inference
+
+
+__all__ = ["PaddleDetectionLayoutModel"]
+
+
+def _resize_image(
+    image: np.ndarray, target_size: Tuple[int, int]
+) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Args:
+        image (np.ndarray): image (np.ndarray)
+    Returns:
+        image (np.ndarray): processed image (np.ndarray)
+    """
+
+    origin_shape = image.shape[:2]
+
+    resize_h, resize_w = target_size
+    # im_scale_y: the resize ratio of Y
+    im_scale_y = resize_h / float(origin_shape[0])
+    # the resize ratio of X
+    im_scale_x = resize_w / float(origin_shape[1])
+
+    # resize image
+    image = cv2.resize(image, None, None, fx=im_scale_x, fy=im_scale_y, interpolation=2)
+    scale_factor = np.array([im_scale_y, im_scale_x]).astype("float32")
+    return image, scale_factor
+
+
+class PaddleDetectionLayoutModel(BaseLayoutModel):
+    """Create a PaddleDetection-based Layout Detection Model
+
+    Args:
+        config_path (:obj:`str`):
+            The path to the configuration file.
+        model_path (:obj:`str`, None):
+            The path to the saved weights of the model.
+            If set, overwrite the weights in the configuration file.
+            Defaults to `None`.
+        label_map (:obj:`dict`, optional):
+            The map from the model prediction (ids) to real
+            word labels (strings). If the config is from one of the supported
+            datasets, Layout Parser will automatically initialize the label_map.
+            Defaults to `None`.
+        device(:obj:`str`, optional):
+            Whether to use cuda or cpu devices. If not set, LayoutParser will
+            automatically determine the device to initialize the models on.
+        extra_config (:obj:`dict`, optional):
+            Extra configuration passed to the PaddleDetection model configuration.
+            Defaults to `{}`.
+            Including arguments:
+            enable_mkldnn (:obj:`bool`, optional):
+                Whether use mkldnn to accelerate the computation.
+                Defaults to False.
+            thread_num (:obj:`int`, optional):
+                The number of threads.
+                Defaults to 10.
+            threshold (:obj:`float`, optional):
+                Threshold to reserve the result for output.
+                Defaults to 0.5.
+            target_size (:obj:`list`, optional):
+                The image shape after resize.
+                Defaults to [640,640].
+
+    Examples::
+        >>> import layoutparser as lp
+        >>> model = lp.models.PaddleDetectionLayoutModel('
+                                    lp://PubLayNet/ppyolov2_r50vd_dcn_365e/config')
+        >>> model.detect(image)
+
+    """
+
+    DEPENDENCIES = ["paddle"]
+    DETECTOR_NAME = "paddledetection"
+    MODEL_CATALOG = MODEL_CATALOG
+
+    def __init__(
+        self,
+        config_path=None,
+        model_path=None,
+        label_map=None,
+        device=None,
+        enforce_cpu=None,
+        extra_config=None,
+    ):
+
+        if enforce_cpu is not None:
+            warnings.warn(
+                "Setting enforce_cpu is deprecated. Please set `device` instead.",
+                DeprecationWarning,
+            )
+
+        if extra_config is None:
+            extra_config = {}
+
+        _, model_path = self.config_parser(config_path, model_path)
+        model_dir = PathManager.get_local_path(model_path)
+
+        if label_map is None:
+            if model_path.startswith("lp://"):
+                dataset_name = model_path.lstrip("lp://").split("/")[1]
+                label_map = LABEL_MAP_CATALOG[dataset_name]
+            else:
+                label_map = {}
+
+        self.label_map = label_map
+
+        # TODO: rethink how to save store the default constants
+        self.predictor = self.load_predictor(
+            model_dir,
+            device=device,
+            enable_mkldnn=extra_config.get("enable_mkldnn", False),
+            thread_num=extra_config.get("thread_num", 10),
+        )
+
+        self.threshold = extra_config.get("threshold", 0.5)
+        self.target_size = extra_config.get("target_size", [640, 640])
+        self.pixel_mean = extra_config.get(
+            "pixel_mean", np.array([[[0.485, 0.456, 0.406]]])
+        )
+        self.pixel_std = extra_config.get(
+            "pixel_std", np.array([[[0.229, 0.224, 0.225]]])
+        )
+
+    def load_predictor(
+        self,
+        model_dir,
+        device=None,
+        enable_mkldnn=False,
+        thread_num=10,
+    ):
+        """set AnalysisConfig, generate AnalysisPredictor
+        Args:
+            model_dir (str): root path of __model__ and __params__
+            device (str): cuda or cpu
+        Returns:
+            predictor (PaddlePredictor): AnalysisPredictor
+        Raises:
+            ValueError: predict by TensorRT need enforce_cpu == False.
+        """
+
+        config = paddle.inference.Config(
+            os.path.join(
+                model_dir, "inference.pdmodel"
+            ),  # TODO: Move them to some constants
+            os.path.join(model_dir, "inference.pdiparams"),
+        )
+
+        if device == "cuda":
+            # initial GPU memory(M), device ID
+            # 2000 is an appropriate value for PaddleDetection model
+            config.enable_use_gpu(2000, 0)
+            # optimize graph and fuse op
+            config.switch_ir_optim(True)
+        else:
+            config.disable_gpu()
+            config.set_cpu_math_library_num_threads(thread_num)
+            if enable_mkldnn:
+                config.enable_mkldnn()
+                try:
+                    # cache 10 different shapes for mkldnn to avoid memory leak
+                    config.set_mkldnn_cache_capacity(10)
+                    config.enable_mkldnn()
+                except Exception as e:
+                    print(
+                        "The current environment does not support `mkldnn`, so disable mkldnn."
+                    )
+
+        # disable print log when predict
+        config.disable_glog_info()
+        # enable shared memory
+        config.enable_memory_optim()
+        # disable feed, fetch OP, needed by zero_copy_run
+        config.switch_use_feed_fetch_ops(False)
+        predictor = paddle.inference.create_predictor(config)
+        return predictor
+
+    def preprocess(self, image):
+        """preprocess image
+
+        Args:
+            image (np.ndarray): image (np.ndarray)
+        Returns:
+            inputs (dict): input of model
+        """
+
+        # resize image by target_size and max_size
+        image, scale_factor = _resize_image(image, self.target_size)
+        input_shape = np.array(image.shape[:2]).astype("float32")
+        # normalize image
+        image = (image / 255.0 - self.pixel_mean) / self.pixel_std
+        # transpose images
+        image = image.transpose((2, 0, 1)).copy()
+
+        inputs = {}
+        inputs["image"] = np.array(image)[np.newaxis, :].astype("float32")
+        inputs["im_shape"] = np.array(input_shape)[np.newaxis, :].astype("float32")
+        inputs["scale_factor"] = np.array(scale_factor)[np.newaxis, :].astype("float32")
+        return inputs
+
+    def gather_output(self, np_boxes):
+        """process output"""
+        layout = Layout()
+        results = []
+        if reduce(lambda x, y: x * y, np_boxes.shape) < 6:
+            print("[WARNING] No object detected.")
+            results = {"boxes": np.array([])}
+        else:
+            results = {}
+            results["boxes"] = np_boxes
+
+        np_boxes = results["boxes"]
+        expect_boxes = (np_boxes[:, 1] > self.threshold) & (np_boxes[:, 0] > -1)
+        np_boxes = np_boxes[expect_boxes, :]
+
+        for np_box in np_boxes:
+            clsid, bbox, score = int(np_box[0]), np_box[2:], np_box[1]
+            x_1, y_1, x_2, y_2 = bbox
+
+            cur_block = TextBlock(
+                Rectangle(x_1, y_1, x_2, y_2),
+                type=self.label_map.get(clsid, clsid),
+                score=score,
+            )
+            layout.append(cur_block)
+
+        return layout
+
+    def detect(self, image):
+        """Detect the layout of a given image.
+
+        Args:
+            image (:obj:`np.ndarray` or `PIL.Image`): The input image to detect.
+
+        Returns:
+            :obj:`~layoutparser.Layout`: The detected layout of the input image
+        """
+
+        # Convert PIL Image Input
+        image = self.image_loader(image)
+
+        inputs = self.preprocess(image)
+
+        input_names = self.predictor.get_input_names()
+
+        for input_name in input_names:
+            input_tensor = self.predictor.get_input_handle(input_name)
+            input_tensor.copy_from_cpu(inputs[input_name])
+
+        self.predictor.run()
+        output_names = self.predictor.get_output_names()
+        boxes_tensor = self.predictor.get_output_handle(output_names[0])
+        np_boxes = boxes_tensor.copy_to_cpu()
+
+        layout = self.gather_output(np_boxes)
+        return layout
+
+    def image_loader(self, image: Union["np.ndarray", "Image.Image"]):
+
+        if isinstance(image, Image.Image):
+            if image.mode != "RGB":
+                image = image.convert("RGB")
+            image = np.array(image)
+
+        return image

layoutparser/ocr/__init__.py

@@ -0,0 +1,16 @@
+# 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 .gcv_agent import GCVAgent, GCVFeatureType
+from .tesseract_agent import TesseractAgent, TesseractFeatureType

layoutparser/ocr/base.py

@@ -0,0 +1,41 @@
+# 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 abc import ABC, abstractmethod
+from enum import IntEnum
+
+from ..file_utils import requires_backends
+
+class BaseOCRElementType(IntEnum):
+    @property
+    @abstractmethod
+    def attr_name(self):
+        pass
+
+
+class BaseOCRAgent(ABC):
+    @property
+    @abstractmethod
+    def DEPENDENCIES(self):
+        """DEPENDENCIES lists all necessary dependencies for the class."""
+        pass
+
+    def __new__(cls, *args, **kwargs):
+
+        requires_backends(cls, cls.DEPENDENCIES)
+        return super().__new__(cls)
+
+    @abstractmethod
+    def detect(self, image):
+        pass

layoutparser/ocr/gcv_agent.py

@@ -0,0 +1,288 @@
+# 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 os
+import json
+import warnings
+
+import numpy as np
+from cv2 import imencode
+
+from .base import BaseOCRAgent, BaseOCRElementType
+from ..elements import Layout, TextBlock, Quadrilateral, TextBlock
+from ..file_utils import is_gcv_available
+
+if is_gcv_available():
+    import google.protobuf.json_format as _json_format
+    import google.cloud.vision as _vision
+
+
+
+def _cvt_GCV_vertices_to_points(vertices):
+    return np.array([[vertex.x, vertex.y] for vertex in vertices])
+
+
+class GCVFeatureType(BaseOCRElementType):
+    """
+    The element types from Google Cloud Vision API
+    """
+
+    PAGE = 0
+    BLOCK = 1
+    PARA = 2
+    WORD = 3
+    SYMBOL = 4
+
+    @property
+    def attr_name(self):
+        name_cvt = {
+            GCVFeatureType.PAGE: "pages",
+            GCVFeatureType.BLOCK: "blocks",
+            GCVFeatureType.PARA: "paragraphs",
+            GCVFeatureType.WORD: "words",
+            GCVFeatureType.SYMBOL: "symbols",
+        }
+        return name_cvt[self]
+
+    @property
+    def child_level(self):
+        child_cvt = {
+            GCVFeatureType.PAGE: GCVFeatureType.BLOCK,
+            GCVFeatureType.BLOCK: GCVFeatureType.PARA,
+            GCVFeatureType.PARA: GCVFeatureType.WORD,
+            GCVFeatureType.WORD: GCVFeatureType.SYMBOL,
+            GCVFeatureType.SYMBOL: None,
+        }
+        return child_cvt[self]
+
+
+class GCVAgent(BaseOCRAgent):
+    """A wrapper for `Google Cloud Vision (GCV) <https://cloud.google.com/vision>`_ Text
+    Detection APIs.
+
+    Note:
+        Google Cloud Vision API returns the output text in two types:
+
+        * `text_annotations`:
+
+            In this format, GCV automatically find the best aggregation
+            level for the text, and return the results in a list. We use
+            :obj:`~gather_text_annotations` to reterive this type of
+            information.
+
+        * `full_text_annotation`:
+
+            To support better user control, GCV also provides the
+            `full_text_annotation` output, where it returns the hierarchical
+            structure of the output text. To process this output, we provide
+            the :obj:`~gather_full_text_annotation` function to aggregate the
+            texts of the given aggregation level.
+    """
+
+    DEPENDENCIES = ["google-cloud-vision"]
+    
+    def __init__(self, languages=None, ocr_image_decode_type=".png"):
+        """Create a Google Cloud Vision OCR Agent.
+
+        Args:
+            languages (:obj:`list`, optional):
+                You can specify the language code of the documents to detect to improve
+                accuracy. The supported language and their code can be found on `this page
+                <https://cloud.google.com/vision/docs/languages>`_.
+                Defaults to None.
+
+            ocr_image_decode_type (:obj:`str`, optional):
+                The format to convert the input image to before sending for GCV OCR.
+                Defaults to `".png"`.
+
+                    * `".png"` is suggested as it does not compress the image.
+                    * But `".jpg"` could also be a good choice if the input image is very large.
+        """
+        try:
+            self._client = _vision.ImageAnnotatorClient()
+        except:
+            warnings.warn(
+                "The GCV credential has not been set. You could not run the detect command."
+            )
+        self._context = _vision.types.ImageContext(language_hints=languages)
+        self.ocr_image_decode_type = ocr_image_decode_type
+
+    @classmethod
+    def with_credential(cls, credential_path, **kwargs):
+        """Specifiy the credential to use for the GCV OCR API.
+
+        Args:
+            credential_path (:obj:`str`): The path to the credential file
+        """
+        os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = credential_path
+        return cls(**kwargs)
+
+    def _detect(self, img_content):
+        img_content = _vision.types.Image(content=img_content)
+        response = self._client.document_text_detection(
+            image=img_content, image_context=self._context
+        )
+        return response
+
+    def detect(
+        self,
+        image,
+        return_response=False,
+        return_only_text=False,
+        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 the google cloud response.
+                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:`~GCVFeatureType`, optional):
+                When set, aggregate the GCV output with respect to the
+                specified aggregation level. Defaults to `None`.
+        """
+        if isinstance(image, np.ndarray):
+            img_content = imencode(self.ocr_image_decode_type, image)[1].tostring()
+
+        elif isinstance(image, str):
+            with io.open(image, "rb") as image_file:
+                img_content = image_file.read()
+
+        res = self._detect(img_content)
+
+        if return_response:
+            return res
+
+        if return_only_text:
+            return res.full_text_annotation.text
+
+        if agg_output_level is not None:
+            return self.gather_full_text_annotation(res, agg_output_level)
+
+        return self.gather_text_annotations(res)
+
+    @staticmethod
+    def gather_text_annotations(response):
+        """Convert the text_annotations from GCV output to an :obj:`Layout` object.
+
+        Args:
+            response (:obj:`AnnotateImageResponse`):
+                The returned Google Cloud Vision AnnotateImageResponse object.
+
+        Returns:
+            :obj:`Layout`: The reterived layout from the response.
+        """
+
+        # The 0th element contains all texts
+        doc = response.text_annotations[1:]
+        gathered_text = Layout()
+
+        for i, text_comp in enumerate(doc):
+            points = _cvt_GCV_vertices_to_points(text_comp.bounding_poly.vertices)
+            gathered_text.append(
+                TextBlock(block=Quadrilateral(points), text=text_comp.description, id=i)
+            )
+
+        return gathered_text
+
+    @staticmethod
+    def gather_full_text_annotation(response, agg_level):
+        """Convert the full_text_annotation from GCV output to an :obj:`Layout` object.
+
+        Args:
+            response (:obj:`AnnotateImageResponse`):
+                The returned Google Cloud Vision AnnotateImageResponse object.
+
+            agg_level (:obj:`~GCVFeatureType`):
+                The layout level to aggregate the text in full_text_annotation.
+
+        Returns:
+            :obj:`Layout`: The reterived layout from the response.
+        """
+
+        def iter_level(
+            iter,
+            agg_level=None,
+            text_blocks=None,
+            texts=None,
+            cur_level=GCVFeatureType.PAGE,
+        ):
+
+            for item in getattr(iter, cur_level.attr_name):
+                if cur_level == agg_level:
+                    texts = []
+
+                # Go down levels to fetch the texts
+                if cur_level == GCVFeatureType.SYMBOL:
+                    texts.append(item.text)
+                elif (
+                    cur_level == GCVFeatureType.WORD
+                    and agg_level != GCVFeatureType.SYMBOL
+                ):
+                    chars = []
+                    iter_level(
+                        item, agg_level, text_blocks, chars, cur_level.child_level
+                    )
+                    texts.append("".join(chars))
+                else:
+                    iter_level(
+                        item, agg_level, text_blocks, texts, cur_level.child_level
+                    )
+
+                if cur_level == agg_level:
+                    nonlocal element_id
+                    points = _cvt_GCV_vertices_to_points(item.bounding_box.vertices)
+                    text_block = TextBlock(
+                        block=Quadrilateral(points),
+                        text=" ".join(texts),
+                        score=item.confidence,
+                        id=element_id,
+                    )
+
+                    text_blocks.append(text_block)
+                    element_id += 1
+
+        if agg_level == GCVFeatureType.PAGE:
+            doc = response.text_annotations[0]
+            points = _cvt_GCV_vertices_to_points(doc.bounding_poly.vertices)
+
+            text_blocks = [TextBlock(block=Quadrilateral(points), text=doc.description)]
+
+        else:
+            doc = response.full_text_annotation
+            text_blocks = []
+            element_id = 0
+            iter_level(doc, agg_level, text_blocks)
+
+        return Layout(text_blocks)
+
+    def load_response(self, filename):
+        with open(filename, "r") as f:
+            data = f.read()
+        return _json_format.Parse(
+            data, _vision.types.AnnotateImageResponse(), ignore_unknown_fields=True
+        )
+
+    def save_response(self, res, file_name):
+        res = _json_format.MessageToJson(res)
+
+        with open(file_name, "w") as f:
+            json_file = json.loads(res)
+            json.dump(json_file, f)