natural-pdf 0.1.38__py3-none-any.whl → 0.2.0__py3-none-any.whl

natural_pdf/core/highlighting_service.py

@@ -6,7 +6,7 @@ import io
 import logging  # Added
 import os
 from dataclasses import dataclass, field
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Literal, Optional, Tuple, Union
 
 from colour import Color
 from PIL import Image, ImageDraw, ImageFont
@@ -17,6 +17,8 @@ try:
 except ImportError:
     Page = Any  # Fallback if circular import issue arises during type checking
 
+from natural_pdf.core.render_spec import RenderSpec
+
 # Import ColorManager and related utils
 from natural_pdf.utils.visualization import (
     ColorManager,
@@ -302,6 +304,134 @@ class HighlightRenderer:
         self.result_image = Image.alpha_composite(self.result_image, overlay)
 
 
+class HighlightContext:
+    """
+    Context manager for accumulating highlights before displaying them together.
+
+    This allows for a clean syntax to show multiple highlight groups:
+
+    Example:
+        with pdf.highlights() as h:
+            h.add(page.find_all('table'), label='tables', color='blue')
+            h.add(page.find_all('text:bold'), label='bold text', color='red')
+            h.show()  # Display all highlights together
+
+    Or for automatic display on exit:
+        with pdf.highlights(show=True) as h:
+            h.add(page.find_all('table'), label='tables')
+            h.add(page.find_all('text:bold'), label='bold')
+            # Automatically shows when exiting the context
+    """
+
+    def __init__(self, source, show_on_exit: bool = False):
+        """
+        Initialize the highlight context.
+
+        Args:
+            source: The source object (PDF, Page, PageCollection, etc.)
+            show_on_exit: If True, automatically show highlights when exiting context
+        """
+        self.source = source
+        self.show_on_exit = show_on_exit
+        self.highlight_groups = []
+        self._color_manager = ColorManager()
+
+    def add(
+        self,
+        elements,
+        label: Optional[str] = None,
+        color: Optional[Union[str, Tuple[int, int, int]]] = None,
+        **kwargs,
+    ) -> "HighlightContext":
+        """
+        Add a group of elements to highlight.
+
+        Args:
+            elements: Elements to highlight (can be ElementCollection, list, or single element)
+            label: Label for this highlight group
+            color: Color for this group (if None, uses color cycling)
+            **kwargs: Additional highlight parameters
+
+        Returns:
+            Self for method chaining
+        """
+        # Convert single element to list
+        if hasattr(elements, "elements"):
+            # It's an ElementCollection
+            element_list = elements.elements
+        elif isinstance(elements, list):
+            element_list = elements
+        else:
+            # Single element
+            element_list = [elements]
+
+        # Determine color if not specified
+        if color is None:
+            color = self._color_manager.get_color(label=label, force_cycle=True)
+
+        self.highlight_groups.append(
+            {"elements": element_list, "label": label, "color": color, **kwargs}
+        )
+
+        return self
+
+    def show(self, **kwargs) -> Optional[Image.Image]:
+        """
+        Display all accumulated highlights.
+
+        Args:
+            **kwargs: Additional parameters passed to the show method
+
+        Returns:
+            PIL Image with all highlights, or None if no source
+        """
+        if not self.source:
+            return None
+
+        # If source has the new unified show method, use it with highlights parameter
+        if hasattr(self.source, "show"):
+            return self.source.show(highlights=self.highlight_groups, **kwargs)
+        else:
+            # Fallback for objects without the new show method
+            logger.warning(
+                f"Object {type(self.source)} does not support unified show() with highlights"
+            )
+            return None
+
+    def render(self, **kwargs) -> Optional[Image.Image]:
+        """
+        Render all accumulated highlights (clean image without debug elements).
+
+        Args:
+            **kwargs: Additional parameters passed to the render method
+
+        Returns:
+            PIL Image with all highlights, or None if no source
+        """
+        if not self.source:
+            return None
+
+        # If source has the new unified render method, use it with highlights parameter
+        if hasattr(self.source, "render"):
+            return self.source.render(highlights=self.highlight_groups, **kwargs)
+        else:
+            # Fallback for objects without the new render method
+            logger.warning(
+                f"Object {type(self.source)} does not support unified render() with highlights"
+            )
+            return None
+
+    def __enter__(self) -> "HighlightContext":
+        """Enter the context."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Exit the context, optionally showing highlights."""
+        if self.show_on_exit and not exc_type:
+            self.show()
+        return False
+
+
 class HighlightingService:
     """
     Central service to manage highlight data and orchestrate rendering.
@@ -418,7 +548,7 @@ class HighlightingService:
         label: Optional[str] = None,
         use_color_cycling: bool = False,
         element: Optional[Any] = None,
-        include_attrs: Optional[List[str]] = None,
+        annotate: Optional[List[str]] = None,
         existing: str = "append",
     ):
         """Adds a rectangular highlight."""
@@ -468,7 +598,7 @@ class HighlightingService:
             label=label,
             use_color_cycling=use_color_cycling,
             element=element,
-            include_attrs=include_attrs,
+            annotate=annotate,
             existing=existing,
         )
 
@@ -480,7 +610,7 @@ class HighlightingService:
         label: Optional[str] = None,
         use_color_cycling: bool = False,
         element: Optional[Any] = None,
-        include_attrs: Optional[List[str]] = None,
+        annotate: Optional[List[str]] = None,
         existing: str = "append",
     ):
         """Adds a polygonal highlight."""
@@ -501,7 +631,7 @@ class HighlightingService:
             label=label,
             use_color_cycling=use_color_cycling,
             element=element,
-            include_attrs=include_attrs,
+            annotate=annotate,
             existing=existing,
         )
 
@@ -514,7 +644,7 @@ class HighlightingService:
         label: Optional[str],
         use_color_cycling: bool,
         element: Optional[Any],
-        include_attrs: Optional[List[str]],
+        annotate: Optional[List[str]],
         existing: str,
     ):
         """Internal method to create and store a Highlight object."""
@@ -533,8 +663,8 @@ class HighlightingService:
 
         # Extract attributes from the element if requested
         attributes_to_draw = {}
-        if element and include_attrs:
-            for attr_name in include_attrs:
+        if element and annotate:
+            for attr_name in annotate:
                 try:
                     attr_value = getattr(element, attr_name, None)
                     if attr_value is not None:
@@ -745,27 +875,59 @@ class HighlightingService:
             else:
                 rendered_image = base_image_pil  # No highlights, no OCR requested
 
-        # --- Add Legend (Based ONLY on this page's highlights) ---
+        # --- Add Legend or Colorbar (Based ONLY on this page's highlights) ---
         if labels:
-            # CHANGE: Create label_colors map only from highlights_on_page
-            labels_colors_on_page: Dict[str, Tuple[int, int, int, int]] = {}
+            # Check if we have quantitative metadata (for colorbar)
+            quantitative_metadata = None
             for hl in highlights_on_page:
-                if hl.label and hl.label not in labels_colors_on_page:
-                    labels_colors_on_page[hl.label] = hl.color
+                if hasattr(hl, "quantitative_metadata") and hl.quantitative_metadata:
+                    quantitative_metadata = hl.quantitative_metadata
+                    break
 
-            if labels_colors_on_page:  # Only add legend if there are labels on this page
-                legend = create_legend(labels_colors_on_page)
-                if legend:  # Ensure create_legend didn't return None
+            if quantitative_metadata:
+                # Create colorbar for quantitative data
+                from natural_pdf.utils.visualization import create_colorbar
+
+                try:
+                    colorbar = create_colorbar(
+                        values=quantitative_metadata["values"],
+                        colormap=quantitative_metadata["colormap"],
+                        bins=quantitative_metadata["bins"],
+                        orientation=(
+                            "horizontal" if legend_position in ["top", "bottom"] else "vertical"
+                        ),
+                    )
                     rendered_image = merge_images_with_legend(
-                        rendered_image, legend, legend_position
+                        rendered_image, colorbar, legend_position
                     )
                     logger.debug(
-                        f"Added legend with {len(labels_colors_on_page)} labels for page {page_index}."
+                        f"Added colorbar for quantitative attribute '{quantitative_metadata['attribute']}' on page {page_index}."
                     )
+                except Exception as e:
+                    logger.warning(f"Failed to create colorbar for page {page_index}: {e}")
+                    # Fall back to regular legend
+                    quantitative_metadata = None
+
+            if not quantitative_metadata:
+                # Create regular categorical legend
+                labels_colors_on_page: Dict[str, Tuple[int, int, int, int]] = {}
+                for hl in highlights_on_page:
+                    if hl.label and hl.label not in labels_colors_on_page:
+                        labels_colors_on_page[hl.label] = hl.color
+
+                if labels_colors_on_page:  # Only add legend if there are labels on this page
+                    legend = create_legend(labels_colors_on_page)
+                    if legend:  # Ensure create_legend didn't return None
+                        rendered_image = merge_images_with_legend(
+                            rendered_image, legend, legend_position
+                        )
+                        logger.debug(
+                            f"Added legend with {len(labels_colors_on_page)} labels for page {page_index}."
+                        )
+                    else:
+                        logger.debug(f"Legend creation returned None for page {page_index}.")
                 else:
-                    logger.debug(f"Legend creation returned None for page {page_index}.")
-            else:
-                logger.debug(f"No labels found on page {page_index}, skipping legend.")
+                    logger.debug(f"No labels found on page {page_index}, skipping legend.")
 
         return rendered_image
 
@@ -875,9 +1037,9 @@ class HighlightingService:
                 )
                 attrs_to_draw = {}
                 element = hl_data.get("element")
-                include_attrs = hl_data.get("include_attrs")
-                if element and include_attrs:
-                    for attr_name in include_attrs:
+                annotate = hl_data.get("annotate")
+                if element and annotate:
+                    for attr_name in annotate:
                         try:
                             attr_value = getattr(element, attr_name, None)
                             if attr_value is not None:
@@ -934,14 +1096,47 @@ class HighlightingService:
 
             legend = None
             if labels:
-                preview_labels = {h.label: h.color for h in preview_highlights if h.label}
-                if preview_labels:
-                    legend = create_legend(preview_labels)
-                    final_image = merge_images_with_legend(
-                        rendered_image, legend, position=legend_position
-                    )
-                else:
-                    final_image = rendered_image
+                # Check if we have quantitative metadata (for colorbar)
+                quantitative_metadata = None
+                for hl_data in temporary_highlights:
+                    if "quantitative_metadata" in hl_data and hl_data["quantitative_metadata"]:
+                        quantitative_metadata = hl_data["quantitative_metadata"]
+                        break
+
+                if quantitative_metadata:
+                    # Create colorbar for quantitative data
+                    from natural_pdf.utils.visualization import create_colorbar
+
+                    try:
+                        colorbar = create_colorbar(
+                            values=quantitative_metadata["values"],
+                            colormap=quantitative_metadata["colormap"],
+                            bins=quantitative_metadata["bins"],
+                            orientation=(
+                                "horizontal" if legend_position in ["top", "bottom"] else "vertical"
+                            ),
+                        )
+                        final_image = merge_images_with_legend(
+                            rendered_image, colorbar, position=legend_position
+                        )
+                        logger.debug(
+                            f"Added colorbar for quantitative attribute '{quantitative_metadata['attribute']}' on page {page_index}."
+                        )
+                    except Exception as e:
+                        logger.warning(f"Failed to create colorbar for page {page_index}: {e}")
+                        # Fall back to regular legend
+                        quantitative_metadata = None
+
+                if not quantitative_metadata:
+                    # Create regular categorical legend
+                    preview_labels = {h.label: h.color for h in preview_highlights if h.label}
+                    if preview_labels:
+                        legend = create_legend(preview_labels)
+                        final_image = merge_images_with_legend(
+                            rendered_image, legend, position=legend_position
+                        )
+                    else:
+                        final_image = rendered_image
             else:
                 final_image = rendered_image
 
@@ -953,3 +1148,519 @@ class HighlightingService:
             raise
 
         return final_image
+
+    def unified_render(
+        self,
+        specs: List["RenderSpec"],
+        resolution: float = 150,
+        width: Optional[int] = None,
+        labels: bool = True,
+        label_format: Optional[str] = None,
+        layout: Literal["stack", "grid", "single"] = "stack",
+        stack_direction: Literal["vertical", "horizontal"] = "vertical",
+        gap: int = 5,
+        columns: Optional[int] = None,
+        background_color: Tuple[int, int, int] = (255, 255, 255),
+        legend_position: str = "right",
+        **kwargs,
+    ) -> Optional[Image.Image]:
+        """
+        Unified rendering method that processes RenderSpec objects.
+
+        This is the single entry point for all image generation in natural-pdf.
+        It handles page rendering, cropping, highlighting, and layout of multiple images.
+
+        Args:
+            specs: List of RenderSpec objects describing what to render
+            resolution: DPI for rendering (default 150)
+            width: Target width in pixels (overrides resolution)
+            labels: Whether to show labels for highlights
+            label_format: Format string for labels
+            layout: How to arrange multiple images
+            stack_direction: Direction for stack layout
+            gap: Pixels between images
+            columns: Number of columns for grid layout
+            background_color: RGB color for background
+            **kwargs: Additional parameters
+
+        Returns:
+            PIL Image or None if nothing to render
+        """
+        from natural_pdf.core.render_spec import RenderSpec
+
+        if not specs:
+            logger.warning("unified_render called with empty specs list")
+            return None
+
+        # Process each spec into an image
+        images = []
+
+        for spec_idx, spec in enumerate(specs):
+            if not isinstance(spec, RenderSpec):
+                logger.error(f"Invalid spec type at index {spec_idx}: {type(spec)}")
+                continue
+
+            try:
+                # Render the page
+                page_image = self._render_spec(
+                    spec=spec,
+                    resolution=resolution,
+                    width=width,
+                    labels=labels,
+                    label_format=label_format,
+                    legend_position=legend_position,
+                    spec_index=spec_idx,
+                    **kwargs,
+                )
+
+                if page_image:
+                    images.append(page_image)
+
+            except Exception as e:
+                logger.error(f"Error rendering spec {spec_idx}: {e}", exc_info=True)
+                continue
+
+        if not images:
+            logger.warning("No images generated from specs")
+            return None
+
+        # Single image - return directly
+        if len(images) == 1:
+            return images[0]
+
+        # Multiple images - apply layout
+        if layout == "stack":
+            return self._stack_images(
+                images, direction=stack_direction, gap=gap, background_color=background_color
+            )
+        elif layout == "grid":
+            return self._grid_images(
+                images, columns=columns, gap=gap, background_color=background_color
+            )
+        else:  # "single" - just return first image
+            logger.warning(f"Multiple specs with layout='single', returning first image only")
+            return images[0]
+
+    def _render_spec(
+        self,
+        spec: "RenderSpec",
+        resolution: float,
+        width: Optional[int],
+        labels: bool,
+        label_format: Optional[str],
+        legend_position: str,
+        spec_index: int,
+        **kwargs,
+    ) -> Optional[Image.Image]:
+        """Render a single RenderSpec to an image."""
+        # Get the page
+        page = spec.page
+        if not hasattr(page, "width") or not hasattr(page, "height"):
+            logger.error(f"Spec {spec_index} page does not have width/height attributes")
+            return None
+
+        # Calculate actual resolution/width
+        if width is not None and page.width > 0:
+            # Calculate resolution from width
+            actual_resolution = (width / page.width) * 72
+        else:
+            # Use provided resolution or default
+            actual_resolution = resolution if resolution is not None else 150
+
+        # Get base page image
+        try:
+            # Use render_plain_page for clean rendering
+            logger.debug(
+                f"Calling render_plain_page with page={page}, resolution={actual_resolution}"
+            )
+            page_image = render_plain_page(page, resolution=actual_resolution)
+        except Exception as e:
+            logger.error(f"Failed to render page: {e}")
+            logger.error(f"Page: {page}, Resolution: {actual_resolution}, Width: {width}")
+            import traceback
+
+            traceback.print_exc()
+            return None
+
+        if page_image is None:
+            return None
+
+        # Apply crop if specified
+        if spec.crop_bbox:
+            page_image = self._crop_image(
+                page_image, spec.crop_bbox, page, actual_resolution / 72  # scale factor
+            )
+
+        # Apply highlights if any
+        if spec.highlights:
+            page_image = self._apply_spec_highlights(
+                page_image,
+                spec.highlights,
+                page,
+                actual_resolution / 72,  # scale factor
+                labels=labels,
+                label_format=label_format,
+                spec_index=spec_index,
+                crop_offset=spec.crop_bbox[:2] if spec.crop_bbox else None,  # Pass crop offset
+            )
+
+            # Add legend or colorbar if labels are enabled
+            if labels:
+                # Import visualization functions
+                from natural_pdf.utils.visualization import (
+                    create_colorbar,
+                    create_legend,
+                    merge_images_with_legend,
+                )
+
+                # Check if we have quantitative metadata (for colorbar)
+                quantitative_metadata = None
+                for highlight_data in spec.highlights:
+                    if (
+                        "quantitative_metadata" in highlight_data
+                        and highlight_data["quantitative_metadata"]
+                    ):
+                        quantitative_metadata = highlight_data["quantitative_metadata"]
+                        break
+
+                if quantitative_metadata:
+                    # Create colorbar for quantitative data
+                    try:
+                        colorbar = create_colorbar(
+                            values=quantitative_metadata["values"],
+                            colormap=quantitative_metadata["colormap"],
+                            bins=quantitative_metadata["bins"],
+                            orientation=(
+                                "horizontal" if legend_position in ["top", "bottom"] else "vertical"
+                            ),
+                        )
+                        page_image = merge_images_with_legend(
+                            page_image, colorbar, position=legend_position
+                        )
+                        logger.debug(
+                            f"Added colorbar for quantitative attribute '{quantitative_metadata['attribute']}' in spec {spec_index}."
+                        )
+                    except Exception as e:
+                        logger.warning(f"Failed to create colorbar for spec {spec_index}: {e}")
+                        # Fall back to regular legend
+                        quantitative_metadata = None
+
+                if not quantitative_metadata:
+                    # Create regular categorical legend
+                    spec_labels = {}
+                    for hl in spec.highlights:
+                        label = hl.get("label")
+                        color = hl.get("color")
+                        if label and color:
+                            # Process color to ensure it's an RGBA tuple
+                            processed_color = self._process_color_input(color)
+                            if processed_color:
+                                spec_labels[label] = processed_color
+                            else:
+                                # Fallback to color manager if processing fails
+                                spec_labels[label] = self._color_manager.get_color(label=label)
+
+                    if spec_labels:
+                        legend = create_legend(spec_labels)
+                        if legend:
+                            page_image = merge_images_with_legend(
+                                page_image, legend, position=legend_position
+                            )
+                            logger.debug(
+                                f"Added legend with {len(spec_labels)} labels for spec {spec_index}."
+                            )
+
+        return page_image
+
+    def _crop_image(
+        self,
+        image: Image.Image,
+        crop_bbox: Tuple[float, float, float, float],
+        page: "Page",
+        scale_factor: float,
+    ) -> Image.Image:
+        """Crop an image to the specified bbox."""
+        # Convert PDF coordinates to pixel coordinates
+        x0, y0, x1, y1 = crop_bbox
+        pixel_bbox = (
+            int(x0 * scale_factor),
+            int(y0 * scale_factor),
+            int(x1 * scale_factor),
+            int(y1 * scale_factor),
+        )
+
+        # Ensure valid crop bounds
+        pixel_bbox = (
+            max(0, pixel_bbox[0]),
+            max(0, pixel_bbox[1]),
+            min(image.width, pixel_bbox[2]),
+            min(image.height, pixel_bbox[3]),
+        )
+
+        if pixel_bbox[2] <= pixel_bbox[0] or pixel_bbox[3] <= pixel_bbox[1]:
+            logger.warning(f"Invalid crop bounds: {crop_bbox}")
+            return image
+
+        return image.crop(pixel_bbox)
+
+    def _apply_spec_highlights(
+        self,
+        image: Image.Image,
+        highlights: List[Dict[str, Any]],
+        page: "Page",
+        scale_factor: float,
+        labels: bool,
+        label_format: Optional[str],
+        spec_index: int,
+        crop_offset: Optional[Tuple[float, float]] = None,
+    ) -> Image.Image:
+        """Apply highlights from a RenderSpec to an image."""
+        # Convert to RGBA for transparency
+        if image.mode != "RGBA":
+            image = image.convert("RGBA")
+
+        # Create overlay for highlights
+        overlay = Image.new("RGBA", image.size, (0, 0, 0, 0))
+        draw = ImageDraw.Draw(overlay)
+
+        # Process each highlight
+        for idx, highlight_dict in enumerate(highlights):
+            # Get geometry
+            bbox = highlight_dict.get("bbox")
+            polygon = highlight_dict.get("polygon")
+
+            if bbox is None and polygon is None:
+                logger.warning(f"Highlight {idx} has no geometry")
+                continue
+
+            # Get color
+            color = highlight_dict.get("color")
+            label = highlight_dict.get("label")
+
+            if color is None:
+                # Use label-based color assignment for consistency
+                color = self._color_manager.get_color(label=label, force_cycle=False)
+            else:
+                # Process color input
+                color = self._process_color_input(color)
+                if color is None:
+                    color = self._color_manager.get_color(label=label, force_cycle=False)
+
+            # Generate label if needed
+            if label is None and labels and label_format:
+                # Generate label from format
+                label = label_format.format(index=idx, spec_index=spec_index, total=len(highlights))
+
+            # Calculate offset for cropped images
+            offset_x = 0
+            offset_y = 0
+            if crop_offset:
+                offset_x = crop_offset[0] * scale_factor
+                offset_y = crop_offset[1] * scale_factor
+
+            # Draw the highlight
+            if polygon:
+                # Scale polygon points and apply offset
+                scaled_polygon = [
+                    (p[0] * scale_factor - offset_x, p[1] * scale_factor - offset_y)
+                    for p in polygon
+                ]
+                draw.polygon(
+                    scaled_polygon, fill=color, outline=(color[0], color[1], color[2], BORDER_ALPHA)
+                )
+            else:
+                # Scale bbox and apply offset
+                x0, y0, x1, y1 = bbox
+                scaled_bbox = [
+                    x0 * scale_factor - offset_x,
+                    y0 * scale_factor - offset_y,
+                    x1 * scale_factor - offset_x,
+                    y1 * scale_factor - offset_y,
+                ]
+                draw.rectangle(
+                    scaled_bbox, fill=color, outline=(color[0], color[1], color[2], BORDER_ALPHA)
+                )
+
+                # Draw attributes if present
+                attributes_to_draw = highlight_dict.get("attributes_to_draw")
+                if attributes_to_draw and scaled_bbox:
+                    self._draw_spec_attributes(draw, attributes_to_draw, scaled_bbox, scale_factor)
+
+        # Composite overlay onto image
+        return Image.alpha_composite(image, overlay)
+
+    def _draw_spec_attributes(
+        self,
+        draw: ImageDraw.Draw,
+        attributes: Dict[str, Any],
+        bbox_scaled: List[float],
+        scale_factor: float,
+    ) -> None:
+        """Draw attribute key-value pairs on the highlight."""
+        try:
+            # Slightly larger font, scaled
+            font_size = max(10, int(8 * scale_factor))
+            # Try to load a font
+            try:
+                font = ImageFont.truetype("Arial.ttf", font_size)
+            except IOError:
+                try:
+                    font = ImageFont.truetype("DejaVuSans.ttf", font_size)
+                except IOError:
+                    font = ImageFont.load_default()
+                    font_size = 10  # Reset size for default font
+        except Exception:
+            font = ImageFont.load_default()
+            font_size = 10
+
+        line_height = font_size + int(4 * scale_factor)  # Scaled line spacing
+        bg_padding = int(3 * scale_factor)
+        max_width = 0
+        text_lines = []
+
+        # Format attribute lines
+        for name, value in attributes.items():
+            if isinstance(value, float):
+                value_str = f"{value:.2f}"  # Format floats
+            else:
+                value_str = str(value)
+            line = f"{name}: {value_str}"
+            text_lines.append(line)
+            try:
+                # Calculate max width for background box
+                max_width = max(max_width, draw.textlength(line, font=font))
+            except AttributeError:
+                # Fallback for older PIL versions
+                bbox = draw.textbbox((0, 0), line, font=font)
+                max_width = max(max_width, bbox[2] - bbox[0])
+
+        if not text_lines:
+            return  # Nothing to draw
+
+        total_height = line_height * len(text_lines)
+
+        # Position near top-right corner with padding
+        x = bbox_scaled[2] - int(2 * scale_factor) - max_width
+        y = bbox_scaled[1] + int(2 * scale_factor)
+
+        # Draw background rectangle (semi-transparent white)
+        bg_x0 = x - bg_padding
+        bg_y0 = y - bg_padding
+        bg_x1 = x + max_width + bg_padding
+        bg_y1 = y + total_height + bg_padding
+        draw.rectangle(
+            [bg_x0, bg_y0, bg_x1, bg_y1],
+            fill=(255, 255, 255, 240),
+            outline=(0, 0, 0, 180),  # Light black outline
+            width=1,
+        )
+
+        # Draw text lines (black)
+        current_y = y
+        for line in text_lines:
+            draw.text((x, current_y), line, fill=(0, 0, 0, 255), font=font)
+            current_y += line_height
+
+    def _stack_images(
+        self,
+        images: List[Image.Image],
+        direction: str,
+        gap: int,
+        background_color: Tuple[int, int, int],
+    ) -> Image.Image:
+        """Stack images vertically or horizontally."""
+        if direction == "vertical":
+            # Calculate dimensions
+            max_width = max(img.width for img in images)
+            total_height = sum(img.height for img in images) + gap * (len(images) - 1)
+
+            # Create canvas
+            canvas = Image.new("RGB", (max_width, total_height), background_color)
+
+            # Paste images
+            y_offset = 0
+            for img in images:
+                # Center horizontally
+                x_offset = (max_width - img.width) // 2
+                # Convert RGBA to RGB if needed
+                if img.mode == "RGBA":
+                    # Create white background
+                    bg = Image.new("RGB", img.size, background_color)
+                    bg.paste(img, mask=img.split()[3])  # Use alpha channel as mask
+                    img = bg
+                canvas.paste(img, (x_offset, y_offset))
+                y_offset += img.height + gap
+
+        else:  # horizontal
+            # Calculate dimensions
+            total_width = sum(img.width for img in images) + gap * (len(images) - 1)
+            max_height = max(img.height for img in images)
+
+            # Create canvas
+            canvas = Image.new("RGB", (total_width, max_height), background_color)
+
+            # Paste images
+            x_offset = 0
+            for img in images:
+                # Center vertically
+                y_offset = (max_height - img.height) // 2
+                # Convert RGBA to RGB if needed
+                if img.mode == "RGBA":
+                    bg = Image.new("RGB", img.size, background_color)
+                    bg.paste(img, mask=img.split()[3])
+                    img = bg
+                canvas.paste(img, (x_offset, y_offset))
+                x_offset += img.width + gap
+
+        return canvas
+
+    def _grid_images(
+        self,
+        images: List[Image.Image],
+        columns: Optional[int],
+        gap: int,
+        background_color: Tuple[int, int, int],
+    ) -> Image.Image:
+        """Arrange images in a grid."""
+        n_images = len(images)
+
+        # Determine grid dimensions
+        if columns is None:
+            # Auto-calculate columns for roughly square grid
+            columns = int(n_images**0.5)
+            if columns * columns < n_images:
+                columns += 1
+
+        rows = (n_images + columns - 1) // columns  # Ceiling division
+
+        # Get max dimensions for cells
+        max_width = max(img.width for img in images)
+        max_height = max(img.height for img in images)
+
+        # Calculate canvas size
+        canvas_width = columns * max_width + (columns - 1) * gap
+        canvas_height = rows * max_height + (rows - 1) * gap
+
+        # Create canvas
+        canvas = Image.new("RGB", (canvas_width, canvas_height), background_color)
+
+        # Place images
+        for idx, img in enumerate(images):
+            row = idx // columns
+            col = idx % columns
+
+            # Calculate position (centered in cell)
+            cell_x = col * (max_width + gap)
+            cell_y = row * (max_height + gap)
+            x_offset = cell_x + (max_width - img.width) // 2
+            y_offset = cell_y + (max_height - img.height) // 2
+
+            # Convert RGBA to RGB if needed
+            if img.mode == "RGBA":
+                bg = Image.new("RGB", img.size, background_color)
+                bg.paste(img, mask=img.split()[3])
+                img = bg
+
+            canvas.paste(img, (x_offset, y_offset))
+
+        return canvas