diff --git a/src/layoutparser/visualization.py b/src/layoutparser/visualization.py index 5f636fe..2c5d9da 100644 --- a/src/layoutparser/visualization.py +++ b/src/layoutparser/visualization.py @@ -1,13 +1,23 @@ -from PIL import Image, ImageFont, ImageDraw -from .elements import * -import numpy as np +from typing import List, Union, Dict, Any, Tuple import functools import os import sys import warnings -import layoutparser 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, + _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") @@ -46,8 +56,8 @@ def _draw_vertical_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("RGB", (text_width, text_height), color=text_background_color) - txt_mask = Image.new("RGB", (text_width, text_height), color=text_background_color) + 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) @@ -81,7 +91,7 @@ def _create_new_canvas(canvas, arrangement, text_background_color): if arrangement == "lr": new_canvas = Image.new( - "RGB", + "RGBA", (canvas.width * 2, canvas.height), color=text_background_color or DEFAULT_TEXT_BACKGROUND, ) @@ -89,7 +99,7 @@ def _create_new_canvas(canvas, arrangement, text_background_color): elif arrangement == "ud": new_canvas = Image.new( - "RGB", + "RGBA", (canvas.width, canvas.height * 2), color=text_background_color or DEFAULT_TEXT_BACKGROUND, ) @@ -108,6 +118,56 @@ def _create_color_palette(types): } +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): @@ -124,11 +184,36 @@ def wrap(canvas, layout, *args, **kwargs): 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, layout, box_width=None, + box_alpha=0, color_map=None, show_element_id=False, show_element_type=False, @@ -136,6 +221,7 @@ def draw_box( 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). @@ -149,6 +235,10 @@ def draw_box( 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 @@ -178,12 +268,26 @@ def draw_box( 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`. """ - draw = ImageDraw.Draw(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) @@ -206,12 +310,9 @@ def draw_box( else color_map.get(ele.type, DEFAULT_OUTLINE_COLOR) ) - if not isinstance(ele, Quadrilateral): - draw.rectangle(ele.coordinates, width=box_width, outline=outline_color) + _draw_box_outline_on_handler(draw, ele, outline_color, box_width) - else: - p = ele.points.ravel().tolist() - draw.line(p + p[:2], width=box_width, fill=outline_color) + _draw_transparent_box_on_handler(draw, ele, outline_color, box_alpha) if show_element_id or show_element_type: text = "" @@ -224,17 +325,23 @@ def draw_box( 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.rectangle( - (start_x, start_y, start_x + text_w, start_y + text_h), - fill=id_text_background_color or DEFAULT_TEXT_BACKGROUND, + + _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 or DEFAULT_TEXT_COLOR, + fill=id_text_color, font=font_obj, ) @@ -250,10 +357,12 @@ def draw_text( 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, ): @@ -289,6 +398,10 @@ def draw_text( `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. @@ -301,10 +414,15 @@ def draw_text( 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. @@ -315,6 +433,14 @@ def draw_text( :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) @@ -327,29 +453,42 @@ def draw_text( text_color = text_color or DEFAULT_TEXT_COLOR text_background_color = text_background_color or DEFAULT_TEXT_BACKGROUND - canvas = _create_new_canvas(canvas, arrangement, text_background_color) - draw = ImageDraw.Draw(canvas) + 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: - p = ( - ele.pad(right=text_box_width, bottom=text_box_width) - .points.ravel() - .tolist() - ) + modified_box = ele.pad(right=text_box_width, bottom=text_box_width) - draw.line(p + p[:2], width=text_box_width, fill=text_box_color) + _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) + 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, text_background_color + ele.text, + font_obj, + text_color, + _get_color_rgb(text_background_color, text_background_alpha), ) if with_box_on_text: diff --git a/tests/test_visualization.py b/tests/test_visualization.py index a670ab2..3e6a242 100644 --- a/tests/test_visualization.py +++ b/tests/test_visualization.py @@ -57,14 +57,18 @@ def test_viz(): font_size=15, text_color="pink", text_background_color="grey", + text_background_alpha=0.1, with_box_on_text=True, text_box_width=2, text_box_color="yellow", + text_box_alpha=0.2, with_layout=True, box_width=1, color_map={None: "blue"}, show_element_id=True, id_font_size=8, + box_alpha=0.25, + id_text_background_alpha=0.25 ) draw_box(image, layout)