Spaces:
Running
Running
File size: 5,007 Bytes
b611f38 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 | """
Bounding Box & Layout Region Visualizer.
Renders color-coded semi-transparent bounding boxes, labels, and confidence tags on PIL images.
"""
import os
import io
import base64
from typing import List, Optional, Tuple, Dict, Any
from PIL import Image, ImageDraw, ImageFont
from core.models import Region, REGION_COLORS, RegionType
def get_default_font(size: int = 12) -> ImageFont.ImageFont:
"""Load default font or TrueType font if available."""
font_candidates = [
"arial.ttf",
"segoeui.ttf",
"tahoma.ttf",
"DejaVuSans.ttf",
"C:\\Windows\\Fonts\\arial.ttf",
"C:\\Windows\\Fonts\\segoeui.ttf"
]
for font_name in font_candidates:
try:
return ImageFont.truetype(font_name, size=size)
except Exception:
continue
return ImageFont.load_default()
def render_annotated_image(
image: Image.Image,
regions: List[Region],
draw_labels: bool = True,
draw_confidence: bool = True,
draw_text_snippet: bool = False,
selected_region_types: Optional[List[str]] = None,
alpha_fill: int = 40
) -> Image.Image:
"""
Renders high-quality color-coded bounding boxes onto the document image.
Uses PIL ImageDraw with an alpha overlay layer for translucent bounding box fills.
"""
# Ensure working on RGBA copy for clean translucency
base_img = image.convert("RGBA")
overlay = Image.new("RGBA", base_img.size, (255, 255, 255, 0))
draw_base = ImageDraw.Draw(base_img)
draw_overlay = ImageDraw.Draw(overlay)
font = get_default_font(size=11)
font_small = get_default_font(size=9)
img_w, img_h = base_img.size
for region in regions:
# Filter if user selected specific region types
if selected_region_types and region.region_type not in selected_region_types:
continue
box = region.box
if not box or len(box) < 4:
continue
x1, y1, x2, y2 = [int(v) for v in box[:4]]
# Clamp to image dimensions
x1 = max(0, min(img_w - 1, x1))
y1 = max(0, min(img_h - 1, y1))
x2 = max(x1 + 1, min(img_w, x2))
y2 = max(y1 + 1, min(img_h, y2))
color_spec = REGION_COLORS.get(region.region_type, REGION_COLORS[RegionType.OTHER.value])
rgb = color_spec["rgb"]
fill_rgba = (rgb[0], rgb[1], rgb[2], alpha_fill)
border_rgba = (rgb[0], rgb[1], rgb[2], 230)
# 1. Draw Translucent Fill on Overlay
draw_overlay.rectangle([x1, y1, x2, y2], fill=fill_rgba)
# 2. Draw Solid Crisp Border on Overlay
line_width = 2 if region.region_type in [RegionType.TABLE.value, RegionType.TITLE_HEADER.value] else 1
draw_overlay.rectangle([x1, y1, x2, y2], outline=border_rgba, width=line_width)
# 3. Draw Category Badge & Label Tag
if draw_labels:
conf_str = f" ({int(region.confidence * 100)}%)" if (draw_confidence and region.confidence is not None) else ""
label_text = f"{region.region_type}{conf_str}"
# Measure label size
try:
bbox_text = font.getbbox(label_text)
tw = bbox_text[2] - bbox_text[0]
th = bbox_text[3] - bbox_text[1]
except Exception:
tw, th = len(label_text) * 6, 12
badge_padding_x = 4
badge_padding_y = 2
badge_w = tw + badge_padding_x * 2
badge_h = th + badge_padding_y * 2
# Position badge above box or just inside if at top edge
badge_y1 = y1 - badge_h if y1 >= badge_h else y1
badge_y2 = badge_y1 + badge_h
badge_x1 = x1
badge_x2 = min(img_w, x1 + badge_w)
# Draw Badge Background (Solid RGB)
draw_overlay.rectangle([badge_x1, badge_y1, badge_x2, badge_y2], fill=(rgb[0], rgb[1], rgb[2], 240))
# Draw Badge Text (White)
draw_overlay.text((badge_x1 + badge_padding_x, badge_y1 + badge_padding_y - 1), label_text, fill=(255, 255, 255, 255), font=font)
# Alpha composite the overlay onto the base image
annotated_rgba = Image.alpha_composite(base_img, overlay)
return annotated_rgba.convert("RGB")
def save_visualization(
image: Image.Image,
regions: List[Region],
output_path: str,
**kwargs
) -> str:
"""Renders and saves the annotated image to disk."""
os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
annotated = render_annotated_image(image, regions, **kwargs)
annotated.save(output_path, "PNG", optimize=True)
return output_path
def image_to_base64_jpeg(image: Image.Image, quality: int = 90) -> str:
"""Converts a PIL image to a base64 encoded data URI."""
buf = io.BytesIO()
image.convert("RGB").save(buf, format="JPEG", quality=quality)
encoded = base64.b64encode(buf.getvalue()).decode("utf-8")
return f"data:image/jpeg;base64,{encoded}"
|