Spaces:
Runtime error
Runtime error
| # demo/visualizer.py | |
| import numpy as np | |
| import cv2 | |
| from typing import Optional | |
| def draw_box_on_frame( | |
| frame: np.ndarray, # [H, W, 3] uint8 RGB | |
| box: list, # [x1, y1, x2, y2] | |
| color: tuple = (255, 255, 0), | |
| label: str = "", | |
| thickness: int = 2, | |
| dashed: bool = False | |
| ) -> np.ndarray: | |
| """Draw a single bounding box on a frame""" | |
| frame = frame.copy() | |
| x1, y1, x2, y2 = [int(v) for v in box] | |
| if dashed: | |
| # Draw dashed rectangle manually | |
| dash_len = 10 | |
| gap_len = 5 | |
| pts = [ | |
| ((x1, y1), (x2, y1)), # top | |
| ((x2, y1), (x2, y2)), # right | |
| ((x2, y2), (x1, y2)), # bottom | |
| ((x1, y2), (x1, y1)), # left | |
| ] | |
| for (px1, py1), (px2, py2) in pts: | |
| dx = px2 - px1 | |
| dy = py2 - py1 | |
| dist = max(abs(dx), abs(dy)) | |
| if dist == 0: | |
| continue | |
| for i in range(0, dist, dash_len + gap_len): | |
| s = i / dist | |
| e = min(i + dash_len, dist) / dist | |
| sx = int(px1 + s * dx) | |
| sy = int(py1 + s * dy) | |
| ex = int(px1 + e * dx) | |
| ey = int(py1 + e * dy) | |
| cv2.line(frame, (sx, sy), (ex, ey), color, thickness) | |
| else: | |
| cv2.rectangle(frame, (x1, y1), (x2, y2), color, thickness) | |
| if label: | |
| cv2.putText( | |
| frame, label, | |
| (x1, max(y1 - 8, 12)), | |
| cv2.FONT_HERSHEY_SIMPLEX, | |
| 0.6, color, 2 | |
| ) | |
| return frame | |
| def draw_trajectory_on_frame( | |
| frame: np.ndarray, | |
| boxes: np.ndarray, # [T, 4] — full trajectory | |
| current_t: int, | |
| color: tuple = (255, 200, 0) | |
| ) -> np.ndarray: | |
| """ | |
| Draw the motion path (center points) up to current frame. | |
| Gives a visual "trail" showing where the object came from. | |
| """ | |
| frame = frame.copy() | |
| centers = np.stack([ | |
| (boxes[:, 0] + boxes[:, 2]) / 2, | |
| (boxes[:, 1] + boxes[:, 3]) / 2 | |
| ], axis=1).astype(int) | |
| # Draw path line | |
| for i in range(1, current_t + 1): | |
| alpha = i / (current_t + 1) # fade older points | |
| c = tuple(int(v * alpha) for v in color) | |
| cv2.line( | |
| frame, | |
| tuple(centers[i-1]), | |
| tuple(centers[i]), | |
| c, 2 | |
| ) | |
| # Draw current center dot | |
| cv2.circle(frame, tuple(centers[current_t]), 5, color, -1) | |
| return frame | |
| def create_comparison_strip( | |
| original: np.ndarray, # [T, H, W, 3] | |
| result: np.ndarray, # [T, H, W, 3] | |
| pred_boxes: np.ndarray, # [T, 4] | |
| sample_ts: list = None # which frames to show | |
| ) -> np.ndarray: | |
| """ | |
| Creates a horizontal strip for visual comparison. | |
| Shows: Original | Result | Diff for N sampled frames. | |
| """ | |
| T = len(original) | |
| if sample_ts is None: | |
| sample_ts = [0, T//4, T//2, 3*T//4, T-1] | |
| rows = [] | |
| for t in sample_ts: | |
| orig_t = original[t].copy() | |
| res_t = result[t].copy() | |
| # Draw box on result | |
| res_t = draw_box_on_frame( | |
| res_t, pred_boxes[t], | |
| color=(0, 255, 0), | |
| label=f"t={t}" | |
| ) | |
| # Amplified diff | |
| diff_t = np.abs( | |
| orig_t.astype(np.int32) - result[t].astype(np.int32) | |
| ) | |
| diff_t = (diff_t * 4).clip(0, 255).astype(np.uint8) | |
| # Add labels | |
| def add_label(img, text): | |
| img = img.copy() | |
| cv2.putText(img, text, (10, 25), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.7, | |
| (255, 255, 255), 2) | |
| cv2.putText(img, text, (10, 25), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.7, | |
| (0, 0, 0), 1) | |
| return img | |
| orig_t = add_label(orig_t, "Original") | |
| res_t = add_label(res_t, "Result") | |
| diff_t = add_label(diff_t, "Diff x4") | |
| row = np.concatenate([orig_t, res_t, diff_t], axis=1) | |
| rows.append(row) | |
| return np.concatenate(rows, axis=0) | |