diff --git a/data_pipeline/droid_dataset_dashboard.py b/data_pipeline/droid_dataset_dashboard.py new file mode 100644 index 000000000..055756a00 --- /dev/null +++ b/data_pipeline/droid_dataset_dashboard.py @@ -0,0 +1,224 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Inspect a track-conditioning dataset: are the clips diverse, and do the +CoTracker tracks actually capture the motion? + +Two concerns this surfaces: + 1. Near-duplicate clips (the frame-0 cluster can over-select the same scene). + 2. Weak control signal: CoTracker uses a frame-0 grid, so anything that ENTERS + the frame later (or that the tracker loses) is never queried -> its motion is + invisible to the tracks. A clip can look dynamic yet have near-static tracks. + +``render`` (CPU; run on a node) writes, per clip: + - __overlay.mp4 the video with the CoTracker tracks drawn (dots + tails), + - __thumb.png first frame, +and a ``stats.json`` with per-clip motion-coverage + a near-duplicate grouping. + +``serve`` (no GPU) is a gradio gallery: clips sorted by motion (least first, so +dead/duplicate clips float to the top), each with its overlay + stats. + + # 1) render artifacts (node) + srun --jobid= --overlap --ntasks=1 .venv/bin/python data_pipeline/droid_dataset_dashboard.py render \ + --data-dir /mnt/weka/home/hao.zhang/shao/data/motion_pipeline/droid_track_200 --out + # 2) serve (login node ok) + .venv/bin/python data_pipeline/droid_dataset_dashboard.py serve --dash-dir --share +""" +from __future__ import annotations + +import argparse +import json +import os +import sys +from pathlib import Path +from typing import Any + +import numpy as np + +sys.path.insert(0, os.path.dirname(__file__)) + +MOVE_THRESH_PX = 8.0 # a point "moves" if it travels > this from frame 0 (input px) +VIS_THRESH = 0.5 + + +def _load_frames(path: str) -> np.ndarray: + try: + from decord import VideoReader, cpu + vr = VideoReader(path, ctx=cpu(0)) + return vr.get_batch(list(range(len(vr)))).asnumpy() + except Exception: # noqa: BLE001 + import imageio.v2 as imageio + rd = imageio.get_reader(path, format="ffmpeg") + fr = np.stack([np.asarray(f) for f in rd], axis=0) + rd.close() + return fr + + +def _clip_stats(tracks: np.ndarray, vis: np.ndarray) -> dict[str, Any]: + """Motion coverage from tracks [T,N,2] (px) + vis [T,N].""" + T, N, _ = tracks.shape + v = vis > VIS_THRESH + # displacement of each point from its frame-0 position, only where visible + disp = np.sqrt(((tracks - tracks[0:1])**2).sum(-1)) # [T,N] + max_disp = np.where(v, disp, 0.0).max(0) # [N] per-point peak travel + moving = max_disp > MOVE_THRESH_PX + return { + "n_points": int(N), + "frac_visible": float(v.mean()), + "frac_moving": float(moving.mean()), # fraction of grid points that ever move + "n_moving": int(moving.sum()), + "mean_motion_px": float(max_disp.mean()), + "p95_motion_px": float(np.percentile(max_disp, 95)), + "max_motion_px": float(max_disp.max()), + } + + +def cmd_render(args: argparse.Namespace) -> None: + from fastvideo.train.callbacks.track_validation import (_draw_overlay, _grid_colors, _subsample) + data_dir = Path(args.data_dir) + manifest = json.loads((data_dir / "videos2caption.json").read_text()) + out = Path(args.out) + out.mkdir(parents=True, exist_ok=True) + import imageio.v2 as imageio + from PIL import Image + + embs: list[np.ndarray] = [] + entries: list[dict[str, Any]] = [] + for rec in manifest: + cid = Path(rec["path"]).stem + vpath = str(data_dir / "videos" / rec["path"]) + ppath = rec.get("points_path") + frames = _load_frames(vpath) + d = np.load(ppath) + tracks = d["tracks"].astype(np.float32)[:frames.shape[0]] # [T,N,2] px + vis = d["visibility"].astype(np.float32)[:frames.shape[0]] + H, W = int(frames.shape[1]), int(frames.shape[2]) + + stats = _clip_stats(tracks, vis) + + # overlay (subsampled grid for clarity) + grid = int(round(tracks.shape[1]**0.5)) + tn = tracks / np.array([W, H], np.float32) # normalize for _subsample + tr_s, vs_s = _subsample(tn, vis, grid, args.stride) + colors = _grid_colors(grid, args.stride)[:tr_s.shape[1]] + trpx = tr_s.copy() + trpx[..., 0] *= W + trpx[..., 1] *= H + ov = _draw_overlay(frames, trpx, vs_s, colors, args.tail, 2, VIS_THRESH) + imageio.mimsave(str(out / f"{cid}__overlay.mp4"), ov, fps=args.fps, macro_block_size=1) + Image.fromarray(frames[0]).save(str(out / f"{cid}__thumb.png")) + + # frame-0 embedding for near-duplicate grouping (32x32 gray, L2-norm) + g = np.asarray(Image.fromarray(frames[0]).convert("L").resize((32, 32)), np.float32).reshape(-1) + embs.append(g / (np.linalg.norm(g) + 1e-8)) + + entries.append({"id": cid, "episode": rec.get("source_episode"), "caption": rec["cap"][0], **stats}) + print( + f"[dash] {cid} move%={stats['frac_moving']*100:4.1f} mean={stats['mean_motion_px']:5.1f}px " + f"max={stats['max_motion_px']:5.1f}px", + flush=True) + + # near-duplicate grouping: union-find on cosine-sim >= threshold + E = np.stack(embs, 0) + sim = E @ E.T + n = len(entries) + parent = list(range(n)) + + def find(i: int) -> int: + while parent[i] != i: + parent[i] = parent[parent[i]] + i = parent[i] + return i + + for i in range(n): + for j in range(i + 1, n): + if sim[i, j] >= args.dup_thresh: + parent[find(i)] = find(j) + groups: dict[int, list[int]] = {} + for i in range(n): + groups.setdefault(find(i), []).append(i) + group_id = {} + for gi, (_, members) in enumerate(sorted(groups.items(), key=lambda kv: -len(kv[1]))): + for m in members: + group_id[m] = gi + for i, e in enumerate(entries): + e["dup_group"] = int(group_id[i]) + + n_groups = len(groups) + n_dead = sum(1 for e in entries if e["frac_moving"] < 0.02) + summary = { + "n_clips": n, + "n_dup_groups": n_groups, + "largest_dup_group": max(len(v) for v in groups.values()), + "n_low_motion_clips(<2%)": n_dead, + "median_frac_moving": float(np.median([e["frac_moving"] for e in entries])), + "median_mean_motion_px": float(np.median([e["mean_motion_px"] for e in entries])), + "dup_thresh": args.dup_thresh, + "move_thresh_px": MOVE_THRESH_PX, + } + (out / "stats.json").write_text(json.dumps({"summary": summary, "clips": entries}, indent=2)) + print(f"[dash] SUMMARY: {json.dumps(summary, indent=2)}", flush=True) + print(f"[dash] wrote dashboard artifacts -> {out}", flush=True) + + +def cmd_serve(args: argparse.Namespace) -> None: + import gradio as gr + dash = Path(args.dash_dir) + data = json.loads((dash / "stats.json").read_text()) + summary, clips = data["summary"], data["clips"] + # sort: least motion first (dead/duplicate clips surface), then by dup group + order = sorted(range(len(clips)), key=lambda i: (clips[i]["frac_moving"], clips[i]["dup_group"])) + + def label(i: int) -> str: + c = clips[i] + return (f"#{i} ep{c['episode']} | move {c['frac_moving']*100:.0f}% | " + f"mean {c['mean_motion_px']:.0f}px | dupgrp {c['dup_group']}") + + gallery_items = [(str(dash / f"{clips[i]['id']}__thumb.png"), label(i)) for i in order] + + def show(evt): # gr.SelectData; unannotated so gradio's get_type_hints doesn't eval it + i = order[evt.index] + c = clips[i] + return (str(dash / f"{c['id']}__overlay.mp4"), json.dumps(c, indent=2)) + + head = (f"### DROID dataset inspection — {summary['n_clips']} clips\n" + f"- near-duplicate groups: **{summary['n_dup_groups']}** " + f"(largest group **{summary['largest_dup_group']}** clips)\n" + f"- low-motion clips (<2% points move): **{summary['n_low_motion_clips(<2%)']}**\n" + f"- median fraction of points moving: **{summary['median_frac_moving']*100:.1f}%**, " + f"median mean-motion **{summary['median_mean_motion_px']:.1f}px**\n\n" + f"Sorted least-motion first. Overlay = CoTracker tracks (dots + tails); " + f"watch for moving objects with NO dots on them (off-frame entries / lost tracks).") + + with gr.Blocks(title="DROID dataset dashboard") as demo: + gr.Markdown(head) + with gr.Row(): + gal = gr.Gallery(value=gallery_items, columns=6, height=560, label="clips (least motion first)") + with gr.Column(): + vid = gr.Video(label="track overlay") + meta = gr.Code(label="clip stats", language="json") + gal.select(show, None, [vid, meta]) + demo.queue().launch(server_name=args.host, server_port=args.port, share=args.share) + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + sub = p.add_subparsers(dest="cmd", required=True) + r = sub.add_parser("render") + r.add_argument("--data-dir", required=True, help="dataset root (videos/ + videos2caption.json + tracks/)") + r.add_argument("--out", required=True) + r.add_argument("--stride", type=int, default=3, help="subsample the NxN grid for overlay clarity") + r.add_argument("--tail", type=int, default=12) + r.add_argument("--fps", type=int, default=12) + r.add_argument("--dup-thresh", type=float, default=0.985, help="cosine-sim on frame-0 to call clips duplicates") + r.set_defaults(func=cmd_render) + s = sub.add_parser("serve") + s.add_argument("--dash-dir", required=True) + s.add_argument("--host", default="0.0.0.0") + s.add_argument("--port", type=int, default=7872) + s.add_argument("--share", action="store_true") + s.set_defaults(func=cmd_serve) + args = p.parse_args() + args.func(args) + + +if __name__ == "__main__": + main() diff --git a/examples/inference/gradio/trackwan/app_action.py b/examples/inference/gradio/trackwan/app_action.py index d72fc3385..992e2b241 100644 --- a/examples/inference/gradio/trackwan/app_action.py +++ b/examples/inference/gradio/trackwan/app_action.py @@ -82,17 +82,29 @@ def swap_checkpoint(name: str) -> str: if STATE.get("cur_ckpt") == name: return f"checkpoint: {name} (loaded)" from safetensors.torch import load_file + try: + from torch.distributed.tensor import DTensor + except Exception: # older torch + from torch.distributed._tensor import DTensor path = STATE["exports"][name] sd = load_file(os.path.join(path, "transformer", "model.safetensors")) tgt = STATE["model"].transformer - p0 = next(tgt.parameters()) - sd = {k: v.to(p0.device, p0.dtype) for k, v in sd.items()} - missing, unexpected = tgt.load_state_dict(sd, strict=False) + # Params are FSDP DTensors (even on 1 GPU); copy into each param's local shard + # instead of load_state_dict (which errors mixing Tensor and DTensor). + params = dict(tgt.named_parameters()) + n_ok = 0 + with torch.no_grad(): + for k, p in params.items(): + if k not in sd: + continue + loc = p.to_local() if isinstance(p, DTensor) else p.data + src = sd[k].to(dtype=loc.dtype, device=loc.device) + if tuple(loc.shape) == tuple(src.shape): + loc.copy_(src) + n_ok += 1 tgt.eval() STATE["cur_ckpt"] = name - warn = "" - if len(missing) > 4 or len(unexpected) > 4: # a few buffers (rope/norm) are expected - warn = f" [WARN missing={len(missing)} unexpected={len(unexpected)}]" + warn = "" if n_ok == len(params) else f" [WARN copied {n_ok}/{len(params)} params]" return f"checkpoint: {name} (swapped){warn}" @@ -139,21 +151,19 @@ def on_ckpt(name): return swap_checkpoint(name) -def on_generate(clip_idx, ckpt_name, traj_json, grid_n, extent, background, steps, seed): +def _build_tracks(traj_json, grid_n, extent, background): + """Parse the drawn trajectory -> (patch, pvis, full_tr, full_vs) or (None, msg).""" import json - clip_idx = int(clip_idx) if not traj_json or traj_json.strip() in ("", "[]"): - return None, None, "Record an action first: move the mouse over the frame and press Space (5 s)." + return None, "Record an action first: move the mouse over the frame and press Space (5 s)." try: traj = json.loads(traj_json) except Exception as e: # noqa: BLE001 - return None, None, f"bad trajectory json: {e}" - swap_checkpoint(ckpt_name) + return None, f"bad trajectory json: {e}" Tpx = STATE["Tpx"] patch, pvis = _patch_tracks(traj, Tpx, int(grid_n), float(extent)) if patch is None: - return None, None, "trajectory too short — hold Space and move the mouse for the full 5 s." - + return None, "trajectory too short — hold Space and move the mouse for the full 5 s." if background == "static (dense)": g = st.make_grid(50) bt, bv = st.static(g, Tpx) @@ -161,31 +171,97 @@ def on_generate(clip_idx, ckpt_name, traj_json, grid_n, extent, background, step full_vs = np.concatenate([bv, pvis], axis=1) else: full_tr, full_vs = patch, pvis + return (patch, pvis, full_tr, full_vs), None + +def _first_frame_cached(clip_idx): + cache = STATE.setdefault("_ff_cache", {}) + if clip_idx not in cache: + cache[clip_idx] = _first_frame_rgb(clip_idx) + return cache[clip_idx] + + +def on_preview(clip_idx, traj_json, grid_n, extent, background): + """Show the synthetic control tracks overlaid on the (frozen) first frame -- BEFORE generating.""" + clip_idx = int(clip_idx) + built, err = _build_tracks(traj_json, grid_n, extent, background) + if built is None: + return None, err + patch, pvis, full_tr, full_vs = built + ff = _first_frame_cached(clip_idx) # [H,W,3] + frames = np.repeat(ff[None], patch.shape[0], axis=0) # freeze frame 0, T copies + out = _overlay_patch(frames, patch, pvis) + path = os.path.join(STATE["out_dir"], f"preview_clip{clip_idx}.mp4") + imageio.mimsave(path, out, fps=24, macro_block_size=1) + nbg = (full_tr.shape[1] - patch.shape[1]) + return path, (f"PREVIEW of the traces to be fed (green = your {int(grid_n)}x{int(grid_n)} arm patch, " + f"{patch.shape[1]} pts + {nbg} static background pts). Press Generate when it looks right.") + + +def on_generate(clip_idx, ckpt_name, traj_json, grid_n, extent, background, steps, seed): + """Generator: streams per-step denoise progress to the log, yields the video at the end.""" + import time + from fastvideo.forward_context import set_forward_context + from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler, ) + clip_idx = int(clip_idx) + built, err = _build_tracks(traj_json, grid_n, extent, background) + if built is None: + yield gr.update(), f"[error] {err}", gr.update() + return + patch, pvis, full_tr, full_vs = built + steps = int(steps) + t0 = time.time() + + yield gr.update(), f"[1/4] swapping to {ckpt_name} ...", gr.update() + swap_checkpoint(ckpt_name) + + model = STATE["model"] + device = model.device + dtype = torch.bfloat16 s = STATE["samples"][clip_idx] - tp = torch.from_numpy(full_tr)[None].float() - tv = torch.from_numpy(full_vs)[None].float() - lat = twi.generate(STATE["model"], first_frame_latent=s["first_frame_latent"], - text_embedding=s["text_embedding"], text_attention_mask=s["text_attention_mask"], - track_points=tp, track_visibility=tv, clip_feature=s["clip_feature"], - num_steps=int(steps), seed=int(seed)) - frames = twi.decode_to_pixels(STATE["model"], lat) + ff = s["first_frame_latent"].to(device, dtype) + cond20 = model._build_i2v_cond_concat(ff) + txt = s["text_embedding"].to(device, dtype) + mask = s["text_attention_mask"].to(device, dtype) + img = s["clip_feature"].to(device, dtype) + tp = torch.from_numpy(full_tr)[None].to(device, dtype) + tv = torch.from_numpy(full_vs)[None].to(device, dtype) + _, _, T, H, W = ff.shape + gen = torch.Generator(device="cpu").manual_seed(int(seed)) + latents = torch.randn((1, 16, T, H, W), generator=gen, dtype=torch.float32).to(device) + sched = FlowMatchEulerDiscreteScheduler(shift=float(model.timestep_shift)) + sched.set_timesteps(steps, device=device) + + for i, t in enumerate(sched.timesteps): + model_in = torch.cat([latents.to(dtype), cond20], dim=1) + ts = t.reshape(1).to(device, dtype) + with torch.no_grad(), torch.autocast(device.type, dtype=dtype), \ + set_forward_context(current_timestep=ts, attn_metadata=None): + v = model.transformer(hidden_states=model_in, encoder_hidden_states=txt, + encoder_attention_mask=mask, timestep=ts, + encoder_hidden_states_image=img, track_points=tp, + track_visibility=tv, return_dict=False) + latents = sched.step(v.float(), t, latents.float(), return_dict=False)[0] + yield gr.update(), f"[2/4] denoising {i + 1}/{steps} ({time.time() - t0:.1f}s)", gr.update() + + yield gr.update(), f"[3/4] decoding latents ... ({time.time() - t0:.1f}s)", gr.update() + frames = twi.decode_to_pixels(model, latents) out = _overlay_patch(frames, patch, pvis) tag = ckpt_name.replace(" ", "") path = os.path.join(STATE["out_dir"], f"action_clip{clip_idx}_{tag}_g{int(grid_n)}.mp4") imageio.mimsave(path, out, fps=24, macro_block_size=1) + yield gr.update(), f"[4/4] computing CoTracker EPE ... ({time.time() - t0:.1f}s)", gr.update() from fastvideo.eval.metrics.motion.cotracker_epe.metric import compute_epe - H, W = frames.shape[1], frames.shape[2] ppx = patch.copy() ppx[..., 0] *= W ppx[..., 1] *= H - res = compute_epe(frames, ppx, pvis, STATE["ct"], STATE["model"].device) + res = compute_epe(frames, ppx, pvis, STATE["ct"], device) epe = res["epe"] - msg = (f"{ckpt_name}: patch EPE = {epe:.2f}px ({res['n_points']} pts followed) — " - f"lower = the {int(grid_n)}x{int(grid_n)} patch tracked your drawn action." - if epe is not None else f"{ckpt_name}: EPE n/a (no points re-tracked)") - return _original_video(clip_idx), path, msg + epe_msg = (f"patch EPE = {epe:.2f}px ({res['n_points']} pts) — lower = the arm patch followed your path." + if epe is not None else "EPE n/a (no points re-tracked)") + yield path, f"[done] generated in {time.time() - t0:.1f}s", epe_msg # ----------------------------------------------------------------------------- ui @@ -241,11 +317,13 @@ def build_ui(num_clips: int, ckpt_names: list[str], Tpx: int): with gr.Row(): steps = gr.Slider(10, 60, value=30, step=5, label="Denoise steps") seed = gr.Number(value=1000, label="Seed") - gen_btn = gr.Button("Generate", variant="primary") - status = gr.Textbox(label="Status", interactive=False) + with gr.Row(): + preview_btn = gr.Button("Preview traces") + gen_btn = gr.Button("Generate", variant="primary") + status = gr.Textbox(label="Status / generation log", interactive=False, lines=2) with gr.Column(scale=1): orig_vid = gr.Video(label="Original clip") - gen_vid = gr.Video(label="Generated (your action overlaid)") + gen_vid = gr.Video(label="Preview traces / Generated (your action overlaid)") epe_box = gr.Textbox(label="Action-following (CoTracker EPE on the patch)", interactive=False) ff_b64 = gr.Textbox(elem_id="ff_b64", visible=False) @@ -256,8 +334,9 @@ def build_ui(num_clips: int, ckpt_names: list[str], Tpx: int): ff_b64.change(None, [ff_b64], None, js="(b64)=>{ window._setbg(b64); }") rec_btn.click(None, None, None, js="()=>{ window._startRec(); }") clr_btn.click(None, None, None, js="()=>{ window._clear(); }") + preview_btn.click(on_preview, [clip, traj_json, grid_n, extent, background], [gen_vid, status]) gen_btn.click(on_generate, [clip, ckpt, traj_json, grid_n, extent, background, steps, seed], - [orig_vid, gen_vid, epe_box]) + [gen_vid, status, epe_box]) demo.load(on_clip, [clip], [caption, ff_b64, orig_vid]) demo.load(None, None, None, js=_canvas_js(Tpx)) return demo