custom-layoutparser 9.9.9__py3-none-any.whl

layoutparser/visualization.py

@@ -0,0 +1,571 @@
+# 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, Optional, Union, Dict, Any, Tuple, Dict
+import functools
+import os
+import sys
+from itertools import cycle
+
+import numpy as np
+from PIL import Image, ImageFont, ImageDraw, ImageColor
+
+import layoutparser
+from .elements import (
+    Layout,
+    Interval,
+    Rectangle,
+    TextBlock,
+    Quadrilateral,
+)
+from .elements.utils import cvt_coordinates_to_points
+
+# We need to fix this ugly hack some time in the future
+_lib_path = os.path.dirname(sys.modules[layoutparser.__package__].__file__)
+_font_path = os.path.join(_lib_path, "misc", "NotoSerifCJKjp-Regular.otf")
+
+DEFAULT_BOX_WIDTH_RATIO = 0.005
+DEFAULT_OUTLINE_COLOR = "red"
+DEAFULT_COLOR_PALETTE = "#f6bd60-#f7ede2-#f5cac3-#84a59d-#f28482"
+# From https://coolors.co/f6bd60-f7ede2-f5cac3-84a59d-f28482
+
+DEFAULT_FONT_PATH = _font_path
+DEFAULT_FONT_SIZE = 15
+DEFAULT_FONT_OBJECT = ImageFont.truetype(DEFAULT_FONT_PATH, DEFAULT_FONT_SIZE)
+DEFAULT_TEXT_COLOR = "black"
+DEFAULT_TEXT_BACKGROUND = "white"
+
+__all__ = ["draw_box", "draw_text"]
+
+
+def _draw_vertical_text(
+    text,
+    image_font,
+    text_color,
+    text_background_color,
+    character_spacing=2,
+    space_width=1,
+):
+    """Helper function to draw text vertically.
+    Ref: https://github.com/Belval/TextRecognitionDataGenerator/blob/7f4c782c33993d2b6f712d01e86a2f342025f2df/trdg/computer_text_generator.py
+    """
+
+    space_height = int(image_font.getsize(" ")[1] * space_width)
+
+    char_heights = [
+        image_font.getsize(c)[1] if c != " " else space_height for c in text
+    ]
+    text_width = max([image_font.getsize(c)[0] for c in text])
+    text_height = sum(char_heights) + character_spacing * len(text)
+
+    txt_img = Image.new("RGBA", (text_width, text_height), color=text_background_color)
+    txt_mask = Image.new("RGBA", (text_width, text_height), color=text_background_color)
+
+    txt_img_draw = ImageDraw.Draw(txt_img)
+    txt_mask_draw = ImageDraw.Draw(txt_mask)
+
+    for i, c in enumerate(text):
+        txt_img_draw.text(
+            (0, sum(char_heights[0:i]) + i * character_spacing),
+            c,
+            fill=text_color,
+            font=image_font,
+        )
+
+    return txt_img.crop(txt_img.getbbox())
+
+
+def _calculate_default_box_width(canvas):
+    return max(1, int(min(canvas.size) * DEFAULT_BOX_WIDTH_RATIO))
+
+
+def _create_font_object(font_size=None, font_path=None):
+
+    if font_size is None and font_path is None:
+        return DEFAULT_FONT_OBJECT
+    else:
+        return ImageFont.truetype(
+            font_path or DEFAULT_FONT_PATH, font_size or DEFAULT_FONT_SIZE
+        )
+
+
+def _create_new_canvas(canvas, arrangement, text_background_color):
+
+    if arrangement == "lr":
+        new_canvas = Image.new(
+            "RGBA",
+            (canvas.width * 2, canvas.height),
+            color=text_background_color or DEFAULT_TEXT_BACKGROUND,
+        )
+        new_canvas.paste(canvas, (canvas.width, 0))
+
+    elif arrangement == "ud":
+        new_canvas = Image.new(
+            "RGBA",
+            (canvas.width, canvas.height * 2),
+            color=text_background_color or DEFAULT_TEXT_BACKGROUND,
+        )
+        new_canvas.paste(canvas, (0, canvas.height))
+
+    else:
+        raise ValueError(f"Invalid direction {arrangement}")
+
+    return new_canvas
+
+
+def _create_color_palette(types):
+    return {
+        type: color
+        for type, color in zip(types, cycle(DEAFULT_COLOR_PALETTE.split("-")))
+    }
+
+
+def _get_color_rgb(color_string: Any, alpha: float) -> Tuple[int, int, int, int]:
+    if color_string[0] == "#" and len(color_string) == 7:
+        # When color string is a hex string
+        color_hex = color_string.lstrip("#")
+        return (
+            *tuple(int(color_hex[i : i + 2], 16) for i in (0, 2, 4)),
+            int(255 * alpha),
+        )
+    else:
+        try:
+            rgb = ImageColor.getrgb(color_string)
+            return rgb + (int(255 * alpha),)
+        except:
+            # ImageColor.getrgb will throw an ValueError when the color is not
+            # a valid color string, even if it is in other formats supported by
+            # PIL. As such, we return the color as it is if the first two cases
+            # are not valid.
+            return color_string
+
+
+def _draw_box_outline_on_handler(draw, block, color, width):
+
+    if not hasattr(block, "points"):
+        points = (cvt_coordinates_to_points(block.coordinates),)
+    else:
+        points = block.points
+
+    vertices = points.ravel().tolist()
+    drawing_vertices = vertices + vertices[:2]
+
+    draw.line(
+        drawing_vertices,
+        width=width,
+        fill=color,
+    )
+
+
+def _draw_transparent_box_on_handler(draw, block, color, alpha):
+
+    if hasattr(block, "points"):
+        vertices = [tuple(block) for block in block.points.tolist()]
+    else:
+        vertices = cvt_coordinates_to_points(block.coordinates)
+
+    draw.polygon(
+        vertices,
+        _get_color_rgb(color, alpha),
+    )
+
+
+def image_loader(func):
+    @functools.wraps(func)
+    def wrap(canvas, layout, *args, **kwargs):
+
+        if isinstance(canvas, Image.Image):
+            if canvas.mode != "RGB":
+                canvas = canvas.convert("RGB")
+            canvas = canvas.copy()
+        elif isinstance(canvas, np.ndarray):
+            canvas = Image.fromarray(canvas)
+        out = func(canvas, layout, *args, **kwargs)
+        return out
+
+    return wrap
+
+
+@image_loader
+def draw_transparent_box(
+    canvas: "Image",
+    blocks: Layout,
+    color_map: Dict = None,
+    alpha: float = 0.25,
+) -> "Image":
+    """Given the image, draw a series of transparent boxes based on the blocks,
+    coloring using the specified color_map.
+    """
+
+    if color_map is None:
+        all_types = set([b.type for b in blocks if hasattr(b, "type")])
+        color_map = _create_color_palette(all_types)
+
+    canvas = canvas.copy()
+    draw = ImageDraw.Draw(canvas, "RGBA")
+
+    for block in blocks:
+        _draw_transparent_box_on_handler(draw, block, color_map[block.type], alpha)
+
+    return canvas
+
+
+@image_loader
+def draw_box(
+    canvas: Image.Image,
+    layout: Layout,
+    box_width: Optional[Union[List[int], int]] = None,
+    box_alpha: Optional[Union[List[float], float]] = None,
+    box_color: Optional[Union[List[str], str]] = None,
+    color_map: Optional[Dict] = None,
+    show_element_id: bool = False,
+    show_element_type: bool = False,
+    id_font_size: Optional[int] = None,
+    id_font_path: Optional[str] = None,
+    id_text_color: Optional[str] = None,
+    id_text_background_color: Optional[str] = None,
+    id_text_background_alpha: Optional[float] = 1,
+):
+    """Draw the layout region on the input canvas(image).
+
+    Args:
+        canvas (:obj:`~np.ndarray` or :obj:`~PIL.Image.Image`):
+            The canvas to draw the layout boxes.
+        layout (:obj:`Layout` or :obj:`list`):
+            The layout of the canvas to show.
+        box_width (:obj:`int` or :obj:`List[int]`, optional):
+            Set to change the width of the drawn layout box boundary.
+            Defaults to None, when the boundary is automatically
+            calculated as the the :const:`DEFAULT_BOX_WIDTH_RATIO`
+            * the maximum of (height, width) of the canvas.
+            If box_with is a list, it will assign different widths to
+            the corresponding layout object, and should have the same
+            length as the number of blocks in `layout`.
+        box_alpha (:obj:`float`  or :obj:`List[float]`, optional):
+            A float or list of floats ranging from 0 to 1. Set to change
+            the alpha of the drawn layout box.
+            Defaults to 0 - the layout box will be fully transparent.
+            If box_alpha is a list of floats, it will assign different
+            alphas to the corresponding layout object, and should have
+            the same length as the number of blocks in `layout`.
+        box_color (:obj:`str`  or :obj:`List[str]`, optional):
+            A string or a list of strings for box colors, e.g.,
+            `['red', 'green', 'blue']` or `'red'`.
+            If box_color is a list of strings, it will assign different
+            colors to the corresponding layout object, and should have
+            the same length as the number of blocks in `layout`.
+            Defaults to None. When `box_color` is set, it will override the
+            `color_map`.
+        color_map (dict, optional):
+            A map from `block.type` to the colors, e.g., `{1: 'red'}`.
+            You can set it to `{}` to use only the
+            :const:`DEFAULT_OUTLINE_COLOR` for the outlines.
+            Defaults to None, when a color palette is is automatically
+            created based on the input layout.
+        show_element_id (bool, optional):
+            Whether to display `block.id` on the top-left corner of
+            the block.
+            Defaults to False.
+        show_element_type (bool, optional):
+            Whether to display `block.type` on the top-left corner of
+            the block.
+            Defaults to False.
+        id_font_size (int, optional):
+            Set to change the font size used for drawing `block.id`.
+            Defaults to None, when the size is set to
+            :const:`DEFAULT_FONT_SIZE`.
+        id_font_path (:obj:`str`, optional):
+            Set to change the font used for drawing `block.id`.
+            Defaults to None, when the :const:`DEFAULT_FONT_OBJECT` is used.
+        id_text_color (:obj:`str`, optional):
+            Set to change the text color used for drawing `block.id`.
+            Defaults to None, when the color is set to
+            :const:`DEFAULT_TEXT_COLOR`.
+        id_text_background_color (:obj:`str`, optional):
+            Set to change the text region background used for drawing `block.id`.
+            Defaults to None, when the color is set to
+            :const:`DEFAULT_TEXT_BACKGROUND`.
+        id_text_background_alpha (:obj:`float`, optional):
+            A float range from 0 to 1. Set to change the alpha of the
+            drawn text.
+            Defaults to 1 - the text box will be solid.
+    Returns:
+        :obj:`PIL.Image.Image`:
+            A Image object containing the `layout` draw upon the input `canvas`.
+    """
+
+    assert 0 <= id_text_background_alpha <= 1, ValueError(
+        f"The id_text_background_alpha value {id_text_background_alpha} is not within range [0,1]."
+    )
+
+    draw = ImageDraw.Draw(canvas, mode="RGBA")
+
+    id_text_background_color = id_text_background_color or DEFAULT_TEXT_BACKGROUND
+    id_text_color = id_text_color or DEFAULT_TEXT_COLOR
+
+    if show_element_id or show_element_type:
+        font_obj = _create_font_object(id_font_size, id_font_path)
+
+    if box_alpha is None:
+        box_alpha = [0] * len(layout)
+    else:
+        if isinstance(box_alpha, (float, int)):
+            box_alpha = [box_alpha] * len(layout)
+
+        if len(box_alpha) != len(layout):
+            raise ValueError(
+                f"The number of alphas {len(box_alpha)} is not equal to the number of blocks {len(layout)}"
+            )
+        if not all(0 <= a <= 1 for a in box_alpha):
+            raise ValueError(
+                f"The box_alpha value {box_alpha} is not within range [0,1]."
+            )
+
+    if box_width is None:
+        box_width = _calculate_default_box_width(canvas)
+        box_width = [box_width] * len(layout)
+    else:
+        if isinstance(box_width, (float, int)):
+            box_width = [box_width] * len(layout)
+
+        if len(box_width) != len(layout):
+            raise ValueError(
+                f"The number of widths {len(box_width)} is not equal to the number of blocks {len(layout)}"
+            )
+
+    if box_color is None:
+        if color_map is None:
+            all_types = set([b.type for b in layout if hasattr(b, "type")])
+            color_map = _create_color_palette(all_types)
+        box_color = [
+            DEFAULT_OUTLINE_COLOR
+            if not isinstance(ele, TextBlock)
+            else color_map.get(ele.type, DEFAULT_OUTLINE_COLOR)
+            for ele in layout
+        ]
+    else:
+        if isinstance(box_color, str):
+            box_color = [box_color] * len(layout)
+
+        if len(box_color) != len(layout):
+            raise ValueError(
+                f"The number of colors {len(box_color)} is not equal to the number of blocks {len(layout)}"
+            )
+
+    # A post check of the lengths of the input lists
+    # To support more versions of python, we do not use 
+    # zip(*, strict=True)
+    assert len(layout) == len(box_color) == len(box_alpha) == len(box_width)
+
+    for idx, (ele, color, alpha, width) in enumerate(
+        zip(layout, box_color, box_alpha, box_width)
+    ):
+
+        if isinstance(ele, Interval):
+            ele = ele.put_on_canvas(canvas)
+
+        if width > 0:
+            _draw_box_outline_on_handler(draw, ele, color, width)
+
+        _draw_transparent_box_on_handler(draw, ele, color, alpha)
+
+        if show_element_id or show_element_type:
+            text = ""
+            if show_element_id:
+                ele_id = ele.id or idx
+                text += str(ele_id)
+            if show_element_type:
+                text = str(ele.type) if not text else text + ": " + str(ele.type)
+
+            start_x, start_y = ele.coordinates[:2]
+            text_w, text_h = font_obj.getsize(text)
+
+            text_box_object = Rectangle(
+                start_x, start_y, start_x + text_w, start_y + text_h
+            )
+            # Add a small background for the text
+
+            _draw_transparent_box_on_handler(
+                draw,
+                text_box_object,
+                id_text_background_color,
+                id_text_background_alpha,
+            )
+
+            # Draw the ids
+            draw.text(
+                (start_x, start_y),
+                text,
+                fill=id_text_color,
+                font=font_obj,
+            )
+
+    return canvas
+
+
+@image_loader
+def draw_text(
+    canvas,
+    layout,
+    arrangement: str = "lr",
+    font_size: Optional[int] = None,
+    font_path: Optional[str] = None,
+    text_color: Optional[str] = None,
+    text_background_color: Optional[str] = None,
+    text_background_alpha: Optional[float] = None,
+    vertical_text: bool = False,
+    with_box_on_text: bool = False,
+    text_box_width: Optional[int] = None,
+    text_box_color: Optional[str] = None,
+    text_box_alpha: Optional[float] = None,
+    with_layout: bool = False,
+    **kwargs,
+):
+    """Draw the (detected) text in the `layout` according to
+    their coordinates next to the input `canvas` (image) for better comparison.
+
+    Args:
+        canvas (:obj:`~np.ndarray` or :obj:`~PIL.Image.Image`):
+            The canvas to draw the layout boxes.
+        layout (:obj:`Layout` or :obj:`list`):
+            The layout of the canvas to show.
+        arrangement (`{'lr', 'ud'}`, optional):
+            The arrangement of the drawn text canvas and the original
+            image canvas:
+            * `lr` - left and right
+            * `ud` - up and down
+            Defaults to 'lr'.
+        font_size (:obj:`str`, optional):
+            Set to change the size of the font used for
+            drawing `block.text`.
+            Defaults to None, when the size is set to
+            :const:`DEFAULT_FONT_SIZE`.
+        font_path (:obj:`str`, optional):
+            Set to change the font used for drawing `block.text`.
+            Defaults to None, when the :const:`DEFAULT_FONT_OBJECT` is used.
+        text_color ([type], optional):
+            Set to change the text color used for drawing `block.text`.
+            Defaults to None, when the color is set to
+            :const:`DEFAULT_TEXT_COLOR`.
+        text_background_color ([type], optional):
+            Set to change the text region background used for drawing
+            `block.text`.
+            Defaults to None, when the color is set to
+            :const:`DEFAULT_TEXT_BACKGROUND`.
+        text_background_alpha (:obj:`float`, optional):
+            A float range from 0 to 1. Set to change the alpha of the
+            background of the canvas.
+            Defaults to 1 - the text box will be solid.
+        vertical_text (bool, optional):
+            Whether the text in a block should be drawn vertically.
+            Defaults to False.
+        with_box_on_text (bool, optional):
+            Whether to draw the layout box boundary of a text region
+            on the text canvas.
+            Defaults to False.
+        text_box_width (:obj:`int`, optional):
+            Set to change the width of the drawn layout box boundary.
+            Defaults to None, when the boundary is automatically
+            calculated as the the :const:`DEFAULT_BOX_WIDTH_RATIO`
+            * the maximum of (height, width) of the canvas.
+        text_box_alpha (:obj:`float`, optional):
+            A float range from 0 to 1. Set to change the alpha of the
+            drawn text box.
+            Defaults to 0 - the text box will be fully transparent.
+        text_box_color (:obj:`int`, optional):
+            Set to change the color of the drawn layout box boundary.
+            Defaults to None, when the color is set to
+            :const:`DEFAULT_OUTLINE_COLOR`.
+        with_layout (bool, optional):
+            Whether to draw the layout boxes on the input (image) canvas.
+            Defaults to False.
+            When set to true, you can pass in the arguments in
+            :obj:`draw_box` to change the style of the drawn layout boxes.
+
+    Returns:
+        :obj:`PIL.Image.Image`:
+            A Image object containing the drawn text from `layout`.
+    """
+
+    if text_background_alpha is None:
+        text_background_alpha = 1
+    if text_box_alpha is None:
+        text_box_alpha = 0
+
+    assert 0 <= text_background_alpha <= 1, ValueError(
+        f"The text_background_color value {text_background_alpha} is not within range [0,1]."
+    )
+    assert 0 <= text_box_alpha <= 1, ValueError(
+        f"The text_box_alpha value {text_box_alpha} is not within range [0,1]."
+    )
+
+    if with_box_on_text:
+        if text_box_width is None:
+            text_box_width = _calculate_default_box_width(canvas)
+
+    if with_layout:
+        canvas = draw_box(canvas, layout, **kwargs)
+
+    font_obj = _create_font_object(font_size, font_path)
+    text_box_color = text_box_color or DEFAULT_OUTLINE_COLOR
+    text_color = text_color or DEFAULT_TEXT_COLOR
+    text_background_color = text_background_color or DEFAULT_TEXT_BACKGROUND
+
+    canvas = _create_new_canvas(
+        canvas,
+        arrangement,
+        _get_color_rgb(text_background_color, text_background_alpha),
+    )
+    draw = ImageDraw.Draw(canvas, "RGBA")
+
+    for idx, ele in enumerate(layout):
+
+        if with_box_on_text:
+            modified_box = ele.pad(right=text_box_width, bottom=text_box_width)
+
+            _draw_box_outline_on_handler(
+                draw, modified_box, text_box_color, text_box_width
+            )
+            _draw_transparent_box_on_handler(
+                draw, modified_box, text_box_color, text_box_alpha
+            )
+
+        if not hasattr(ele, "text") or ele.text == "":
+            continue
+
+        (start_x, start_y) = ele.coordinates[:2]
+        if not vertical_text:
+            draw.text(
+                (start_x, start_y),
+                ele.text,
+                font=font_obj,
+                fill=text_color,
+            )
+        else:
+            text_segment = _draw_vertical_text(
+                ele.text,
+                font_obj,
+                text_color,
+                _get_color_rgb(text_background_color, text_background_alpha),
+            )
+
+            if with_box_on_text:
+                # Avoid cover the box regions
+                canvas.paste(
+                    text_segment, (start_x + text_box_width, start_y + text_box_width)
+                )
+            else:
+                canvas.paste(text_segment, (start_x, start_y))
+
+    return canvas