Skip to content

Commit

Permalink
Better Visualization Functions (#129)
Browse files Browse the repository at this point in the history 
* Not drawing outlines when box_width is 0

* Support colors in draw_box

* Better naming & variable annotation

* support seeting box_alpha & box_width as a list
  • Loading branch information
@lolipopshock
lolipopshock authored Apr 3, 2022
1 parent e7c62d6 commit 895530f
Show file tree
Hide file tree
Showing 2 changed files with 145 additions and 66 deletions.
145 changes: 98 additions & 47 deletions src/layoutparser/visualization.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import List, Union, Dict, Any, Tuple
from typing import List, Optional, Union, Dict, Any, Tuple, Dict
import functools
import os
import sys
Expand Down Expand Up @@ -223,18 +223,19 @@ def draw_transparent_box(

@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,
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).
Expand All @@ -243,15 +244,29 @@ def draw_box(
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):
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.
box_alpha (:obj:`float`, optional):
A float range from 0 to 1. Set to change the alpha of the
drawn layout box.
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
Expand Down Expand Up @@ -290,9 +305,6 @@ def draw_box(
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]."
)
Expand All @@ -302,30 +314,66 @@ def draw_box(
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)
if box_alpha is None:
box_alpha = [0]
else:
if isinstance(box_alpha, (float, int)):
box_alpha = [box_alpha] * len(layout)

for idx, ele in enumerate(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 isinstance(ele, Interval):
ele = ele.put_on_canvas(canvas)
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)}"
)

outline_color = (
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)

_draw_box_outline_on_handler(draw, ele, outline_color, box_width)
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)}"
)

_draw_transparent_box_on_handler(draw, ele, outline_color, box_alpha)
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 = ""
Expand Down Expand Up @@ -365,18 +413,18 @@ def draw_box(
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,
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
Expand All @@ -392,7 +440,6 @@ def draw_text(
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
Expand Down Expand Up @@ -435,7 +482,6 @@ def draw_text(
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.
Expand All @@ -447,6 +493,11 @@ def draw_text(
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]."
)
Expand Down
66 changes: 47 additions & 19 deletions tests/test_visualization.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.

import pytest

from layoutparser.elements import *
from layoutparser.ocr import *
from layoutparser.visualization import *
Expand All @@ -30,26 +32,52 @@ def test_viz():
draw_box(image, Layout([]))
draw_text(image, Layout([]))

draw_box(
image,
Layout(
[
Interval(0, 10, axis="x"),
Rectangle(0, 50, 100, 80),
Quadrilateral(np.array([[10, 10], [30, 40], [90, 40], [10, 20]])),
]
),
layout = Layout(
[
Interval(0, 10, axis="x"),
Rectangle(0, 50, 100, 80),
Quadrilateral(np.array([[10, 10], [30, 40], [90, 40], [10, 20]])),
]
)

draw_text(
draw_box(image, layout)
draw_text(image, layout)

# Test colors
draw_box(image, layout, box_color=["red", "green", "blue"])
draw_box(image, layout, box_color="red")

draw_text(image, layout, box_color=["red", "green", "blue"])
with pytest.raises(ValueError):
draw_box(image, layout, box_color=["red", "green", "blue", "yellow"])
with pytest.raises(ValueError):
draw_text(
image,
layout,
box_color=["red", "green", "blue", "yellow"],
with_layout=True,
)

# Test alphas
draw_box(image, layout, box_alpha=0)
draw_box(image, layout, box_alpha=[0.1, 0.2, 0.3])
with pytest.raises(ValueError):
draw_box(image, layout, box_color=[0.1, 0.2, 0.3, 0.5])
with pytest.raises(ValueError):
draw_box(image, layout, box_color=[0.1, 0.2, 0.3, 1.5])

# Test widths
draw_box(image, layout, box_width=1)
draw_box(image, layout, box_width=[1, 2, 3])
with pytest.raises(ValueError):
draw_box(image, layout, box_width=[1, 2, 3, 4])

draw_box(
image,
Layout(
[
Interval(0, 10, axis="x"),
Rectangle(0, 50, 100, 80),
Quadrilateral(np.array([[10, 10], [30, 40], [90, 40], [10, 20]])),
]
),
layout,
box_alpha=[0.1, 0.2, 0.3],
box_width=[1, 2, 3],
box_color=["red", "green", "blue"],
)

for idx, level in enumerate(
Expand Down Expand Up @@ -82,8 +110,8 @@ def test_viz():
show_element_id=True,
id_font_size=8,
box_alpha=0.25,
id_text_background_alpha=0.25
id_text_background_alpha=0.25,
)

draw_box(image, layout)
draw_text(image, layout)
draw_text(image, layout)

0 comments on commit 895530f

Please sign in to comment.