Spaces:
Runtime error
Runtime error
File size: 2,834 Bytes
f3d0a26 | 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 | # demo/preview.py
import numpy as np
from visualizer import draw_box_on_frame, draw_trajectory_on_frame
from utils.box_utils import interpolate_boxes
def preview_trajectory(
first_frame: np.ndarray, # [H, W, 3]
start_box: list, # [x1, y1, x2, y2]
end_box: list, # [x1, y1, x2, y2]
num_frames: int = 81
) -> np.ndarray:
"""
Shows the planned trajectory on the first frame BEFORE running.
User sees this immediately after drawing boxes — fast feedback.
"""
keyboxes = {0: start_box, num_frames - 1: end_box}
boxes = interpolate_boxes(keyboxes, num_frames)
frame = first_frame.copy()
# Draw full trajectory path (center points)
centers = np.stack([
(boxes[:, 0] + boxes[:, 2]) / 2,
(boxes[:, 1] + boxes[:, 3]) / 2
], axis=1).astype(int)
for i in range(1, len(centers)):
alpha = i / len(centers)
color = (
int(255 * (1 - alpha)),
int(200 * alpha),
255
)
import cv2
cv2.line(frame,
tuple(centers[i-1]),
tuple(centers[i]),
color, 2)
# Draw start box (solid yellow)
frame = draw_box_on_frame(
frame, start_box,
color=(255, 220, 0),
label="START",
dashed=False
)
# Draw end box (dashed yellow)
frame = draw_box_on_frame(
frame, end_box,
color=(255, 220, 0),
label="END",
dashed=True
)
# Draw a few intermediate boxes (faded)
for i in [20, 40, 60]:
if i < len(boxes):
frame = draw_box_on_frame(
frame, boxes[i],
color=(200, 200, 200),
label=f"t={i}",
dashed=True,
thickness=1
)
return frame
def preview_trajectory_strip(
frames: np.ndarray, # [T, H, W, 3]
start_box: list,
end_box: list,
) -> np.ndarray:
"""
Shows predicted box overlaid on 5 sampled frames.
Gives sense of how box moves through the video.
"""
T = len(frames)
keyboxes = {0: start_box, T - 1: end_box}
boxes = interpolate_boxes(keyboxes, T)
sample_ts = [0, T//4, T//2, 3*T//4, T-1]
previews = []
for t in sample_ts:
frame = frames[t].copy()
frame = draw_box_on_frame(
frame, boxes[t],
color=(0, 255, 255),
label=f"t={t}",
dashed=(t > 0)
)
# Add small frame counter
import cv2
H, W = frame.shape[:2]
progress = f"{t}/{T-1}"
cv2.putText(frame, progress, (W-80, H-10),
cv2.FONT_HERSHEY_SIMPLEX, 0.5,
(200, 200, 200), 1)
previews.append(frame)
return np.concatenate(previews, axis=1) # horizontal strip
|