trajectory_viz.py

"""
Trajectory Visualization Utilities for Trace Model

Extracts trajectory coordinates from model output text and overlays them on images.
Supports both pixel coordinates and normalized (0-1) coordinates.
"""

import os
import re
from typing import List, Tuple, Optional, Union

import numpy as np
from PIL import Image, ImageDraw


def extract_trajectory_from_text(text: str) -> List[List[float]]:
    """
    Extract trajectory coordinates from model output text.

    Handles both pixel coordinates [[100, 200], [150, 250]] and
    normalized coordinates [[0.5, 0.3], [0.7, 0.4]].

    Args:
        text: The text output from the model containing trajectory information

    Returns:
        List of [x, y] coordinate pairs as floats
    """
    # Look for coordinate pairs [x, y] - supports ints and floats
    coord_pattern = r"\[\s*(-?\d+(?:\.\d+)?)\s*,\s*(-?\d+(?:\.\d+)?)\s*\]"
    coord_matches = re.findall(coord_pattern, text)

    if not coord_matches:
        return []

    trajectory = []
    for x_str, y_str in coord_matches:
        try:
            x = float(x_str.strip())
            y = float(y_str.strip())
            trajectory.append([x, y])
        except (ValueError, IndexError):
            continue

    return trajectory


def _to_pixel_coords(
    trajectory: List[List[float]],
    img_width: int,
    img_height: int,
    normalized: bool = True,
) -> List[List[int]]:
    """Convert trajectory to pixel coordinates."""
    pixel_traj = []
    for x, y in trajectory:
        if normalized:
            x = int(x * img_width)
            y = int(y * img_height)
        else:
            x, y = int(x), int(y)
        pixel_traj.append([x, y])
    return pixel_traj


def visualize_trajectory_on_image(
    trajectory: List[List[float]],
    image_path: Optional[str] = None,
    output_path: Optional[str] = None,
    pil_image: Optional[Image.Image] = None,
    normalized: bool = True,
    start_color: Tuple[int, int, int] = (0, 255, 0),
    end_color: Tuple[int, int, int] = (255, 0, 0),
    line_thickness: int = 4,
) -> Optional[np.ndarray]:
    """
    Overlay trajectory on an image with gradient coloring (green start -> red end).

    Args:
        trajectory: List of [x, y] coordinate pairs (pixel or normalized)
        image_path: Path to input image (used if pil_image is None)
        output_path: Where to save the output image
        pil_image: PIL Image to draw on (overrides image_path)
        normalized: If True, coordinates are 0-1 and will be scaled to image size
        start_color: RGB for trajectory start
        end_color: RGB for trajectory end
        line_thickness: Line width in pixels

    Returns:
        numpy array of the output image, or None if trajectory too short
    """
    if not trajectory or len(trajectory) < 2:
        return None

    if pil_image is not None:
        img = pil_image.convert("RGB").copy()
    elif image_path and os.path.exists(image_path):
        img = Image.open(image_path).convert("RGB").copy()
    else:
        return None

    w, h = img.size
    pixel_traj = _to_pixel_coords(trajectory, w, h, normalized=normalized)

    # Clamp to image bounds
    pixel_traj = [
        [max(0, min(w - 1, x)), max(0, min(h - 1, y))]
        for x, y in pixel_traj
    ]

    draw = ImageDraw.Draw(img)

    # Draw gradient line segments
    num_segments = len(pixel_traj) - 1
    for i in range(num_segments):
        progress = i / max(1, num_segments - 1)
        r = int(start_color[0] * (1 - progress) + end_color[0] * progress)
        g = int(start_color[1] * (1 - progress) + end_color[1] * progress)
        b = int(start_color[2] * (1 - progress) + end_color[2] * progress)
        segment_color = (r, g, b)
        start_pt = tuple(pixel_traj[i])
        end_pt = tuple(pixel_traj[i + 1])
        draw.line([start_pt, end_pt], fill=segment_color, width=line_thickness)

    # Draw start marker
    if pixel_traj:
        sx, sy = pixel_traj[0]
        r = max(3, line_thickness)
        bbox = [sx - r, sy - r, sx + r, sy + r]
        draw.ellipse(bbox, fill=start_color, outline=(255, 255, 255), width=2)

    if output_path:
        img.save(output_path)

    return np.array(img)