Source code for layoutparser.visualization

# 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, Dict, Any, Tuple
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


[docs]@image_loader def draw_box( canvas, layout, box_width=None, box_alpha=0, color_map=None, show_element_id=False, show_element_type=False, id_font_size=None, id_font_path=None, id_text_color=None, id_text_background_color=None, id_text_background_alpha=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`, 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. box_alpha (:obj:`float`, optional): A float range from 0 to 1. Set to change the alpha of the drawn layout box. Defaults to 0 - the layout box will be fully transparent. 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 <= box_alpha <= 1, ValueError( f"The box_alpha value {box_alpha} is not within range [0,1]." ) 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 box_width is None: box_width = _calculate_default_box_width(canvas) if show_element_id or show_element_type: font_obj = _create_font_object(id_font_size, id_font_path) 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) for idx, ele in enumerate(layout): if isinstance(ele, Interval): ele = ele.put_on_canvas(canvas) outline_color = ( DEFAULT_OUTLINE_COLOR if not isinstance(ele, TextBlock) else color_map.get(ele.type, DEFAULT_OUTLINE_COLOR) ) _draw_box_outline_on_handler(draw, ele, outline_color, box_width) _draw_transparent_box_on_handler(draw, ele, outline_color, box_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
[docs]@image_loader def draw_text( canvas, layout, arrangement="lr", font_size=None, font_path=None, text_color=None, text_background_color=None, text_background_alpha=1, vertical_text=False, with_box_on_text=False, text_box_width=None, text_box_color=None, text_box_alpha=0, with_layout=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`. """ 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