File size: 4,185 Bytes
23e1f91 | 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 | """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
# Trails (semi-transparent)
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])
# draw segments where both endpoints visible
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)
# Current points: filled circle if visible, hollow ring otherwise
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)
# Side-by-side panel of key frames
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()
|