| """Visualize a tracks.npz on top of source frames: point IDs are color-coded, |
| visibility shown via marker style, and trails up to the current frame are drawn.""" |
| import argparse, os |
| import numpy as np |
| import cv2 |
| import imageio |
|
|
|
|
| def load_video(path, max_frames=None): |
| cap = cv2.VideoCapture(path); frames=[] |
| while True: |
| r, f = cap.read() |
| if not r: break |
| frames.append(cv2.cvtColor(f, cv2.COLOR_BGR2RGB)) |
| if max_frames and len(frames) >= max_frames: break |
| cap.release(); return frames |
|
|
|
|
| def colormap(N): |
| rng = np.random.default_rng(0) |
| return (rng.uniform(0, 255, (N, 3))).astype(np.uint8) |
|
|
|
|
| def draw_frame(frame, tracks_so_far, vis_so_far, colors, trail_alpha=0.6, radius=4): |
| """tracks_so_far: (t+1, N, 2). Draw trail and current-pt for each id.""" |
| img = frame.copy() |
| overlay = img.copy() |
| T_now, N, _ = tracks_so_far.shape |
|
|
| |
| for n in range(N): |
| pts = tracks_so_far[:, n, :].astype(np.int32) |
| v = vis_so_far[:, n] |
| col = tuple(int(c) for c in colors[n]) |
| |
| for i in range(1, T_now): |
| if v[i] and v[i-1]: |
| cv2.line(overlay, tuple(pts[i-1]), tuple(pts[i]), col, 1, cv2.LINE_AA) |
| img = cv2.addWeighted(overlay, trail_alpha, img, 1 - trail_alpha, 0) |
|
|
| |
| cur = tracks_so_far[-1].astype(np.int32) |
| cur_v = vis_so_far[-1] |
| for n in range(N): |
| col = tuple(int(c) for c in colors[n]) |
| x, y = int(cur[n, 0]), int(cur[n, 1]) |
| if cur_v[n]: |
| cv2.circle(img, (x, y), radius, col, -1, cv2.LINE_AA) |
| cv2.circle(img, (x, y), radius, (0, 0, 0), 1, cv2.LINE_AA) |
| else: |
| cv2.circle(img, (x, y), radius, col, 1, cv2.LINE_AA) |
| return img |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--video", required=True) |
| ap.add_argument("--tracks_npz", required=True) |
| ap.add_argument("--out_dir", required=True) |
| ap.add_argument("--max_frames", type=int, default=0) |
| ap.add_argument("--keyframes", type=int, nargs="+", default=None, |
| help="frames to dump as PNG (defaults to evenly spaced 4)") |
| ap.add_argument("--make_video", action="store_true", |
| help="also dump a track-with-trails MP4") |
| args = ap.parse_args() |
|
|
| os.makedirs(args.out_dir, exist_ok=True) |
| frames = load_video(args.video, args.max_frames or None) |
| z = np.load(args.tracks_npz) |
| tracks, vis = z["tracks"], z["vis"].astype(bool) |
| T = min(len(frames), tracks.shape[0]) |
| frames, tracks, vis = frames[:T], tracks[:T], vis[:T] |
| N = tracks.shape[1] |
| colors = colormap(N) |
| print(f"T={T}, N={N}") |
| print(f"vis frame 0: {int(vis[0].sum())}/{N}") |
| print(f"vis frame -1: {int(vis[-1].sum())}/{N}") |
|
|
| keyframes = args.keyframes or list(np.linspace(0, T - 1, 4).astype(int)) |
| panels = [] |
| for k in keyframes: |
| img = draw_frame(frames[k], tracks[:k+1], vis[:k+1], colors) |
| cv2.putText(img, f"frame {k} vis={int(vis[k].sum())}/{N}", (16, 36), |
| cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 255, 255), 3, cv2.LINE_AA) |
| cv2.putText(img, f"frame {k} vis={int(vis[k].sum())}/{N}", (16, 36), |
| cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 0), 1, cv2.LINE_AA) |
| out_p = os.path.join(args.out_dir, f"trail_frame{k:03d}.png") |
| imageio.imwrite(out_p, img) |
| print(f" saved {out_p}") |
| panels.append(img) |
|
|
| |
| H, W, _ = panels[0].shape |
| panel = np.concatenate(panels, axis=1) |
| panel_p = os.path.join(args.out_dir, "panel.png") |
| imageio.imwrite(panel_p, panel) |
| print(f" saved {panel_p} ({panel.shape[1]}x{panel.shape[0]})") |
|
|
| if args.make_video: |
| out = [] |
| for t in range(T): |
| out.append(draw_frame(frames[t], tracks[:t+1], vis[:t+1], colors)) |
| outp = os.path.join(args.out_dir, "tracks_with_trails.mp4") |
| imageio.mimwrite(outp, out, fps=15, macro_block_size=1) |
| print(f" saved {outp}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
