[feat] WanTrack: dataset-inspection dashboard + interactive draw-trace upgrades
- droid_dataset_dashboard.py: render per-clip CoTracker overlays + motion-coverage stats (frac_moving, mean/max motion) + near-duplicate grouping, and a gradio gallery sorted least-motion-first to surface weak/duplicate clips. - app_action.py: DTensor-safe checkpoint hot-swap (copy into each param's local shard instead of load_state_dict, which errors on FSDP DTensors); a Preview-traces button that overlays the synthetic control on the frozen first frame before generating; and a streaming per-step denoise log.
This commit is contained in:
shaoxiongduan
@@ -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:
|
||||
- <id>__overlay.mp4 the video with the CoTracker tracks drawn (dots + tails),
|
||||
- <id>__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=<job> --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 <dash_dir>
|
||||
# 2) serve (login node ok)
|
||||
.venv/bin/python data_pipeline/droid_dataset_dashboard.py serve --dash-dir <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()
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user