Spaces:
Sleeping
Sleeping
File size: 4,199 Bytes
7c21061 | 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 | """
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)
|