Spaces:
Running
Running
| import torch | |
| import cv2 | |
| import numpy as np | |
| from skimage.color import gray2rgb | |
| def points_to_tensor(points: list, qt: int, orig_H: int, orig_W: int, target: int = 256) -> torch.Tensor: | |
| """ | |
| Convert [(x1,y1), ..., (xn,yn)] to tensor of shape [1, n, 3] | |
| where last dim is (qt, x, y), with x/y scaled to target resolution. | |
| Args: | |
| points : list of (x, y) tuples or np.array([x, y]) | |
| qt : single int, same for all points | |
| orig_H : original frame height | |
| orig_W : original frame width | |
| target : target resolution (default 256) | |
| Returns: | |
| tensor of shape [1, n, 3], dtype float32 | |
| """ | |
| scale_x = target / orig_W | |
| scale_y = target / orig_H | |
| arr = np.array( | |
| [[qt, p[0] * scale_x, p[1] * scale_y] for p in points], | |
| dtype=np.float32 | |
| ) # (n, 3) | |
| return torch.tensor(arr).unsqueeze(0) # (1, n, 3) | |
| def visualize_tracking( | |
| frames: np.ndarray, | |
| points: np.ndarray, | |
| tracking_quality: np.ndarray = None, | |
| vis_color='random', | |
| color_map: np.ndarray = None, | |
| gray: bool = False, | |
| alpha: float = 1.0, | |
| track_length: int = 0, | |
| thickness: int = 2, | |
| ) -> np.ndarray: | |
| num_points, num_frames = points.shape[:2] | |
| height, width = frames.shape[1:3] | |
| if gray and frames.shape[-1] != 3: | |
| frames = gray2rgb(frames.squeeze()) | |
| radius = max(6, int(0.006 * min(height, width))) | |
| quality_colors = { | |
| 0: np.array([255, 0, 0]), | |
| 1: np.array([255, 255, 0]), | |
| 2: np.array([0, 255, 0]), | |
| } | |
| video = frames.copy() | |
| # Stable random colors | |
| if vis_color == 'random' and tracking_quality is None and color_map is None: | |
| rand_colors = np.random.randint(0, 256, size=(num_points, 3)) | |
| for t in range(num_frames): | |
| overlay = np.zeros_like(video[t], dtype=np.uint8) | |
| t_start = max(1, t - track_length) | |
| for i in range(num_points): | |
| # ------------------------------------------------- | |
| # Resolve color ONCE (fixes UnboundLocalError) | |
| # ------------------------------------------------- | |
| if tracking_quality is not None: | |
| color = quality_colors.get( | |
| int(tracking_quality[i, t]), | |
| np.array([255, 255, 255]) | |
| ) | |
| elif color_map is not None: | |
| color = np.asarray(color_map[i]) | |
| elif isinstance(vis_color, (list, tuple, np.ndarray)): | |
| color = np.asarray(vis_color) | |
| else: | |
| if vis_color == 'random': | |
| color = rand_colors[i] | |
| elif vis_color == 'red': | |
| color = quality_colors[0] | |
| elif vis_color == 'yellow': | |
| color = quality_colors[1] | |
| elif vis_color == 'green': | |
| color = quality_colors[2] | |
| else: | |
| raise ValueError(f"Unknown vis_color: {vis_color}") | |
| color = color.astype(np.uint8) | |
| # ------------------------------------------------- | |
| # Draw track lines | |
| # ------------------------------------------------- | |
| for tt in range(t_start, t): | |
| fade = (tt - t_start + 1) / max(1, (t - t_start)) | |
| x0n, y0n = points[i, tt - 1] | |
| x1n, y1n = points[i, tt] | |
| x0 = int(np.clip(x0n * width, 0, width - 1)) | |
| y0 = int(np.clip(y0n * height, 0, height - 1)) | |
| x1 = int(np.clip(x1n * width, 0, width - 1)) | |
| y1 = int(np.clip(y1n * height, 0, height - 1)) | |
| faded_color = (color * fade).astype(np.uint8) | |
| cv2.line( | |
| overlay, | |
| (x0, y0), | |
| (x1, y1), | |
| faded_color.tolist(), | |
| thickness=thickness, | |
| lineType=cv2.LINE_AA | |
| ) | |
| # ------------------------------------------------- | |
| # Draw dot (current position) | |
| # ------------------------------------------------- | |
| xc = int(points[i, t, 0] * width) | |
| yc = int(points[i, t, 1] * height) | |
| cv2.circle( | |
| overlay, | |
| (xc, yc), | |
| radius=radius, | |
| color=color.tolist(), | |
| thickness=-1 | |
| ) | |
| video[t] = cv2.addWeighted(video[t], 1.0, overlay, alpha, 0) | |
| return video |