track_video.py

#!/usr/bin/env python3
"""Example: track points through a video with one of the wrapped trackers.

    python track_video.py --video clip.mp4 --method tapnext
    python track_video.py --video frames_dir/ --method cow_tapnext --n-points 600 --query-frame 0 --out results/

By default the queries are a regular grid over the WHOLE first frame (--n-points or --spacing). Use --mask to restrict
them to a binary mask image. Writes <out>/<name>_<method>.npz (tracks [N,T,2] x,y pixels, visible [N,T], queries [N,3]
(t,x,y)) and an mp4 with the track trails (--no-viz to skip)."""
import argparse
import time
from pathlib import Path

import cv2
import numpy as np

from tracking import METHODS, draw_tracks, grid_queries, load_video


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--video", required=True, help="video file (mp4, ...) or a directory of frames")
    ap.add_argument("--method", default="tapnext", choices=sorted(METHODS))
    ap.add_argument("--out", default="results")
    ap.add_argument("--n-points", type=int, default=400, help="approximate number of grid queries")
    ap.add_argument("--spacing", type=int, default=None, help="grid spacing in pixels (overrides --n-points)")
    ap.add_argument("--query-frame", type=int, default=0)
    ap.add_argument("--mask", default=None, help="optional binary mask image: only query inside it")
    ap.add_argument("--max-frames", type=int, default=None)
    ap.add_argument("--stride", type=int, default=1, help="use every k-th frame")
    ap.add_argument("--max-side", type=int, default=None, help="downscale so the long side is at most this")
    ap.add_argument("--no-viz", action="store_true")
    a = ap.parse_args()

    video = load_video(a.video, a.max_frames, a.stride, a.max_side)
    T, H, W, _ = video.shape
    mask = None
    if a.mask:
        mask = cv2.resize(cv2.imread(a.mask, cv2.IMREAD_GRAYSCALE), (W, H), interpolation=cv2.INTER_NEAREST) > 127
    queries = grid_queries(H, W, a.spacing, a.n_points, t=a.query_frame, mask=mask)
    print(f"video {T} frames {W}x{H}, {len(queries)} queries at frame {a.query_frame}, method {a.method}")

    t0 = time.time()
    res = METHODS[a.method](video, queries)
    print(f"{res.summary()}  ({time.time() - t0:.1f} s)")

    out = Path(a.out); out.mkdir(parents=True, exist_ok=True)
    stem = Path(a.video.rstrip("/")).stem
    res.save(out / f"{stem}_{a.method}.npz")
    if not a.no_viz:
        draw_tracks(video, res, str(out / f"{stem}_{a.method}.mp4"))
    print(f"wrote {out / f'{stem}_{a.method}.npz'}")


if __name__ == "__main__":
    main()