core/visualizer.py

"""
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}"