Spaces:
Running on Zero
Running on Zero
| import spaces # must be imported before torch / anything touching CUDA | |
| import os | |
| import sys | |
| import time | |
| import tempfile | |
| from functools import lru_cache | |
| import numpy as np | |
| import torch | |
| import gradio as gr | |
| import trimesh | |
| import imageio.v2 as imageio | |
| from PIL import Image, ImageDraw, ImageFont | |
| from huggingface_hub import hf_hub_download | |
| ROOT = os.path.dirname(os.path.abspath(__file__)) | |
| sys.path.insert(0, ROOT) | |
| sys.path.insert(0, os.path.join(ROOT, "src")) | |
| from src.train import CUT3RDynamicsModel, _adapt_state_dict_for_height_input # noqa: E402 | |
| from src.data.multi_dataloader import MultiObjectDynamics # noqa: E402 | |
| MODEL_REPO = "frankzydou/RigidFormer" | |
| DATA_REPO = "frankzydou/RigidFormer-MOVi" | |
| DATASETS = { | |
| # UI name: (parquet config, checkpoint) | |
| "MOVi-B": ("movi_b", "rigidformer_movib.pth"), | |
| "MOVi-A": ("movi_a", "rigidformer_movia.pth"), | |
| "MOVi-Sphere": ("movi_sphere", "rigidformer_movisphere.pth"), | |
| } | |
| SHAPES = ("cone", "cube", "cylinder", "gear", "sphere", "sponge", "spot", "suzanne", "teapot", "torus", "torus_knot") | |
| FPS_PHYS = 240 # physical frames per second of the MOVi simulations | |
| WARMUP = 4 # ground-truth context frames (evaluation protocol of the paper) | |
| ANCHOR_K = 4 | |
| VIDEO_FPS = 48 # one rendered frame every 5 physical frames (real time) | |
| GRAVITY = np.array([0.0, 0.0, -10.0]) | |
| # -------------------------------------------------------------------------------------- | |
| # Data: test scenes of RigidFormer-MOVi (CC BY 4.0) and KuBasic collision meshes (CC BY-SA 4.0) | |
| # -------------------------------------------------------------------------------------- | |
| COLLISION = {} | |
| for s in SHAPES: | |
| p = hf_hub_download(DATA_REPO, f"assets/kubasic/{s}/collision_geometry.obj", repo_type="dataset") | |
| # file order of the vertices == the order of the object's vertices in the training data | |
| m = trimesh.load(p, process=False) | |
| COLLISION[s] = (np.asarray(m.vertices, np.float64), np.asarray(m.faces, np.int64)) | |
| def _render_mesh(shape): | |
| V, F = COLLISION[shape] | |
| if shape == "sphere": # nicer sphere for display only (same radius as the collision sphere) | |
| ico = trimesh.creation.icosphere(subdivisions=3, radius=0.5) | |
| V, F = np.asarray(ico.vertices), np.asarray(ico.faces) | |
| return V, F | |
| RENDER_MESH = {s: _render_mesh(s) for s in SHAPES} | |
| TRAJ_COLS = ("positions", "quaternions", "velocities") | |
| def _read_rows(path): | |
| import pyarrow.parquet as pq | |
| cols = ["scene", "num_objects", "shape", "size", "material", "mass", "friction", "restitution", "color", | |
| "positions", "quaternions", "velocities"] | |
| table = pq.read_table(path, columns=cols) | |
| frame_cols = [(k, 3 if k != "quaternions" else 4) for k in TRAJ_COLS] | |
| plain = table.drop_columns([k for k, _ in frame_cols]).to_pylist() | |
| for key, d in frame_cols: | |
| col = table.column(key).combine_chunks() | |
| frame_off = col.offsets.to_numpy() | |
| frames = col.flatten() | |
| obj_off = frames.offsets.to_numpy() | |
| values = frames.flatten().flatten().to_numpy(zero_copy_only=False).reshape(-1, d) | |
| for r, rec in enumerate(plain): | |
| f0, f1 = frame_off[r], frame_off[r + 1] | |
| rec[key] = values[obj_off[f0]:obj_off[f1]].reshape(f1 - f0, rec["num_objects"], d).astype(np.float64) | |
| return plain | |
| SCENES = {} | |
| for name, (cfg, _) in DATASETS.items(): | |
| p = hf_hub_download(DATA_REPO, f"data/{cfg}/test-00000-of-00001.parquet", repo_type="dataset") | |
| SCENES[name] = _read_rows(p) | |
| print(f"[data] {name}: {len(SCENES[name])} test scenes", flush=True) | |
| # -------------------------------------------------------------------------------------- | |
| # Models (one per dataset), loaded at module scope and moved to CUDA eagerly (ZeroGPU) | |
| # -------------------------------------------------------------------------------------- | |
| MODELS = {} | |
| for name, (_, ckpt_name) in DATASETS.items(): | |
| ckpt = torch.load(hf_hub_download(MODEL_REPO, ckpt_name), map_location="cpu", weights_only=False) | |
| config = ckpt["config"] | |
| model = CUT3RDynamicsModel(config) | |
| sd = {k[len("module."):] if k.startswith("module.") else k: v for k, v in ckpt["model_state_dict"].items()} | |
| sd = _adapt_state_dict_for_height_input(model, sd) | |
| model.load_state_dict(sd, strict=True) | |
| model.eval() | |
| MODELS[name] = model.to("cuda") | |
| del ckpt, sd | |
| print(f"[model] {name}: {ckpt_name} loaded", flush=True) | |
| # -------------------------------------------------------------------------------------- | |
| # Geometry helpers | |
| # -------------------------------------------------------------------------------------- | |
| def quat_wxyz_to_rotmat(q): | |
| q = np.asarray(q, np.float64) | |
| q = q / np.linalg.norm(q, axis=-1, keepdims=True) | |
| w, x, y, z = q[..., 0], q[..., 1], q[..., 2], q[..., 3] | |
| R = np.stack([ | |
| 1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w), | |
| 2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w), | |
| 2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)], axis=-1) | |
| return R.reshape(q.shape[:-1] + (3, 3)) | |
| def kabsch_poses(reference, frames): | |
| """Rigid poses (R (T,3,3), t (T,3)) taking reference (V,3) to frames (T,V,3).""" | |
| X0, X = np.asarray(reference, np.float64), np.asarray(frames, np.float64) | |
| c0, c = X0.mean(0), X.mean(1) | |
| U, _, Vt = np.linalg.svd(np.einsum("vi,tvj->tij", X0 - c0, X - c[:, None])) | |
| V, Ut = np.transpose(Vt, (0, 2, 1)), np.transpose(U, (0, 2, 1)) | |
| D = np.tile(np.eye(3), (len(X), 1, 1)) | |
| D[:, 2, 2] = np.sign(np.linalg.det(V @ Ut)) | |
| R = V @ D @ Ut | |
| return R, c - np.einsum("tij,j->ti", R, c0) | |
| def object_pool(dataset): | |
| """Every object of the dataset's test scenes as it appears at frame 0 (for random new scenes).""" | |
| pool = [] | |
| for rec in SCENES[dataset]: | |
| for i in range(rec["num_objects"]): | |
| pool.append({k: rec[k][i] for k in ("shape", "size", "material", "mass", "friction", "restitution", "color")} | |
| | {"quat": rec["quaternions"][0, i]}) | |
| return pool | |
| def random_scene(dataset, num_objects, seed, speed=1.0, horizon=480): | |
| """A new MOVi-style scene: objects drawn from the dataset's object pool, dropped from 1.5-4.5 m with a | |
| horizontal throw toward the centre (as in the MOVi generator). Only the first context frames are known | |
| (free flight before any contact), so the result has no ground truth.""" | |
| rng = np.random.default_rng(int(seed)) | |
| pool = object_pool(dataset) | |
| objs, placed = [], [] | |
| tries = 0 | |
| while len(objs) < num_objects and tries < 5000: | |
| tries += 1 | |
| o = pool[rng.integers(len(pool))] | |
| radius = float(o["size"]) * np.linalg.norm(COLLISION[o["shape"]][0], axis=1).max() | |
| pos = np.array([rng.uniform(-4.0, 4.0), rng.uniform(-4.0, 4.0), rng.uniform(1.5, 4.5)]) | |
| if any(np.linalg.norm(pos - q) < radius + r + 0.05 for q, r in placed): | |
| continue | |
| placed.append((pos, radius)) | |
| target = rng.uniform(-4.0, 4.0, size=2) | |
| vel = np.array([*(speed * (target - pos[:2])), 0.0]) | |
| objs.append(dict(o, pos=pos, vel=vel)) | |
| n = len(objs) | |
| T = horizon | |
| t = np.arange(T)[:, None] / FPS_PHYS | |
| positions = np.stack([o["pos"][None] + o["vel"][None] * t + 0.5 * GRAVITY[None] * t ** 2 for o in objs], 1) | |
| quats = np.stack([np.tile(o["quat"], (T, 1)) for o in objs], 1) | |
| return { | |
| "scene": f"random (seed {int(seed)})", "num_objects": n, | |
| "shape": [o["shape"] for o in objs], "size": [float(o["size"]) for o in objs], | |
| "material": [o["material"] for o in objs], "mass": [float(o["mass"]) for o in objs], | |
| "friction": [float(o["friction"]) for o in objs], "restitution": [float(o["restitution"]) for o in objs], | |
| "color": [o["color"] for o in objs], "positions": positions, "quaternions": quats, | |
| "ground_truth": False, | |
| } | |
| def _fps_cached(key, V_bytes, n): | |
| V = np.frombuffer(V_bytes, np.float32).reshape(n, 3) | |
| anchors = MultiObjectDynamics._farthest_point_sample_np(V, ANCHOR_K, rule="centroid_farthest") | |
| order = MultiObjectDynamics._farthest_point_sample_np(V, n, rule="centroid_farthest") | |
| return anchors, order | |
| def build_batch(rec, step_size): | |
| """The model input exactly as the authors' data loader builds it for evaluation (step size s: context | |
| frames 0, s, 2s, 3s; reference pose = global frame 0; FPS anchors and vertex ranks from the stored, | |
| unscaled frame-0 vertices).""" | |
| n = rec["num_objects"] | |
| nvs = [len(COLLISION[s][0]) for s in rec["shape"]] | |
| max_nv = max(nvs) | |
| frames = [k * step_size for k in range(WARMUP)] | |
| R_all = quat_wxyz_to_rotmat(rec["quaternions"][: frames[-1] + 1]) # (T, n, 3, 3) | |
| P_all = rec["positions"][: frames[-1] + 1] | |
| def verts(t, i, scaled=True): | |
| v = COLLISION[rec["shape"][i]][0] | |
| s = rec["size"][i] if scaled else 1.0 | |
| return P_all[t, i] + s * v @ R_all[t, i].T | |
| obj = np.zeros((1, WARMUP, n, max_nv, 3), np.float32) | |
| vmask = np.zeros((1, WARMUP, n, max_nv), bool) | |
| ref = np.zeros((1, n, max_nv, 3), np.float32) | |
| anchor_idx = np.zeros((1, n, ANCHOR_K), np.int64) | |
| vtx_rank = np.full((1, n, max_nv), max_nv, np.int64) | |
| for i in range(n): | |
| nv = nvs[i] | |
| for k, t in enumerate(frames): | |
| obj[0, k, i, :nv] = verts(t, i) | |
| vmask[0, k, i, :nv] = True | |
| ref[0, i, :nv] = verts(0, i) | |
| V0 = verts(0, i, scaled=False).astype(np.float32) | |
| anchors, order = _fps_cached(rec["shape"][i], V0.tobytes(), nv) | |
| anchor_idx[0, i] = anchors | |
| vtx_rank[0, i, order] = np.arange(nv) | |
| phys = np.array([[rec["mass"][i], rec["friction"][i], rec["restitution"][i]] for i in range(n)], np.float32)[None] | |
| return { | |
| "obj_positions": obj, "valid_obj_mask": np.ones((1, n), bool), "valid_vtx_mask": vmask, | |
| "ref_obj_positions": ref, "anchor_idx": anchor_idx, "vtx_rank": vtx_rank, | |
| "phys_params": phys, "skip_step_size": np.array([float(step_size)], np.float32), | |
| } | |
| def run_rollout(dataset, batch, num_steps): | |
| model = MODELS[dataset] | |
| tb = {k: torch.from_numpy(v).to("cuda") for k, v in batch.items()} | |
| torch.cuda.synchronize() | |
| t0 = time.time() | |
| try: | |
| # the authors' --cuda-graph fast path: one step is captured and replayed (bit-identical trajectories) | |
| os.environ["RIGIDFORMER_CUDA_GRAPH"] = "1" | |
| out = model.rollout(initial_batch=tb, num_steps=num_steps, dt=1.0, return_extras=False, warmup_frames=WARMUP) | |
| except Exception as e: # fall back to the default eager path | |
| print(f"[rollout] CUDA-graph path failed ({type(e).__name__}: {e}); using the eager path", flush=True) | |
| os.environ["RIGIDFORMER_CUDA_GRAPH"] = "0" | |
| tb = {k: torch.from_numpy(v).to("cuda") for k, v in batch.items()} | |
| out = model.rollout(initial_batch=tb, num_steps=num_steps, dt=1.0, return_extras=False, warmup_frames=WARMUP) | |
| torch.cuda.synchronize() | |
| elapsed = time.time() - t0 | |
| print(f"[rollout] {dataset}: {num_steps} frames, step size {float(batch['skip_step_size'][0]):g}, " | |
| f"{batch['obj_positions'].shape[2]} objects: {elapsed:.2f} s", flush=True) | |
| return out["obj_positions"][0].float().cpu().numpy(), elapsed | |
| # -------------------------------------------------------------------------------------- | |
| # Rendering (CPU, painter's algorithm with PIL) | |
| # -------------------------------------------------------------------------------------- | |
| W, H, SS = 640, 480, 2 | |
| LIGHT = np.array([0.35, -0.45, 0.82]); LIGHT /= np.linalg.norm(LIGHT) | |
| def _camera(): | |
| target = np.array([0.0, 0.0, 0.6]) | |
| eye = np.array([0.0, -11.5, 7.8]) | |
| f = target - eye; f /= np.linalg.norm(f) | |
| r = np.cross(f, [0, 0, 1.0]); r /= np.linalg.norm(r) | |
| u = np.cross(r, f) | |
| focal = (H * SS / 2) / np.tan(np.radians(42) / 2) | |
| return eye, f, r, u, focal | |
| CAM = _camera() | |
| def _project(X): | |
| eye, f, r, u, focal = CAM | |
| d = X - eye | |
| z = d @ f | |
| x = d @ r | |
| y = d @ u | |
| z = np.maximum(z, 1e-3) | |
| return np.stack([W * SS / 2 + focal * x / z, H * SS / 2 - focal * y / z], -1), z | |
| def _floor_image(): | |
| img = Image.new("RGB", (W * SS, H * SS), (236, 238, 242)) | |
| draw = ImageDraw.Draw(img) | |
| # horizon gradient | |
| for yy in range(0, H * SS, 4): | |
| c = int(246 - 18 * yy / (H * SS)) | |
| draw.rectangle([0, yy, W * SS, yy + 4], fill=(c, c, c + 4)) | |
| quads = [] | |
| for i in range(-10, 10): | |
| for j in range(-10, 10): | |
| Q = np.array([[i, j, 0], [i + 1, j, 0], [i + 1, j + 1, 0], [i, j + 1, 0]], float) | |
| xy, z = _project(Q) | |
| if (z < 0.5).any(): | |
| continue | |
| shade = 222 if (i + j) % 2 == 0 else 200 | |
| quads.append((z.mean(), xy, shade)) | |
| for _, xy, shade in sorted(quads, key=lambda q: -q[0]): | |
| draw.polygon([tuple(p) for p in xy], fill=(shade, shade, shade + 6)) | |
| return img | |
| def render_frame(objects, title, subtitle=""): | |
| """objects: list of (vertices (V,3) world, faces (F,3), rgb 0-1).""" | |
| img = _floor_image().copy() | |
| draw = ImageDraw.Draw(img, "RGBA") | |
| eye = CAM[0] | |
| # blob shadows | |
| for V, _, _ in objects: | |
| c = V.mean(0) | |
| rad = 0.5 * (V.max(0) - V.min(0))[:2].max() | |
| h = max(V[:, 2].min(), 0.0) | |
| a = int(110 * np.exp(-h / 2.0)) | |
| ring = np.array([[c[0] + rad * np.cos(t), c[1] + rad * np.sin(t), 0.001] for t in np.linspace(0, 2 * np.pi, 24)]) | |
| xy, z = _project(ring) | |
| if (z > 0.5).all(): | |
| draw.polygon([tuple(p) for p in xy], fill=(40, 40, 50, a)) | |
| tris = [] | |
| for V, F, rgb in objects: | |
| A, B, C = V[F[:, 0]], V[F[:, 1]], V[F[:, 2]] | |
| N = np.cross(B - A, C - A) | |
| N /= np.linalg.norm(N, axis=1, keepdims=True) + 1e-12 | |
| centers = (A + B + C) / 3 | |
| vis = ((eye - centers) * N).sum(1) > 0 | |
| if not vis.any(): | |
| continue | |
| xy, z = _project(V) | |
| lam = np.clip(N[vis] @ LIGHT, 0, 1) | |
| view = eye - centers[vis]; view /= np.linalg.norm(view, axis=1, keepdims=True) | |
| hvec = view + LIGHT; hvec /= np.linalg.norm(hvec, axis=1, keepdims=True) | |
| spec = np.clip((N[vis] * hvec).sum(1), 0, 1) ** 24 | |
| base = np.asarray(rgb, float) | |
| col = np.clip(base[None] * (0.32 + 0.68 * lam[:, None]) + 0.35 * spec[:, None], 0, 1) | |
| col = (col * 255).astype(int) | |
| depth = z[F[vis]].mean(1) | |
| for d, tri, c in zip(depth, F[vis], col): | |
| tris.append((d, xy[tri], tuple(c))) | |
| tris.sort(key=lambda q: -q[0]) | |
| for _, xy, c in tris: | |
| draw.polygon([tuple(p) for p in xy], fill=c) | |
| img = img.resize((W, H), Image.LANCZOS) | |
| d2 = ImageDraw.Draw(img) | |
| try: | |
| font = ImageFont.load_default(size=22) | |
| small = ImageFont.load_default(size=15) | |
| except TypeError: | |
| font = small = ImageFont.load_default() | |
| d2.text((14, 10), title, fill=(30, 30, 40), font=font) | |
| if subtitle: | |
| d2.text((14, 38), subtitle, fill=(70, 70, 85), font=small) | |
| return np.asarray(img) | |
| def _color(c): | |
| c = np.clip(np.asarray(c, float), 0, 1) | |
| return 0.15 + 0.8 * c # keep colours vivid but never pure black/white | |
| def poses_to_objects(rec, R, t): | |
| """R (n,3,3), t (n,3) -> renderable objects.""" | |
| objs = [] | |
| for i in range(rec["num_objects"]): | |
| V, F = RENDER_MESH[rec["shape"][i]] | |
| Vw = t[i] + rec["size"][i] * V @ R[i].T | |
| objs.append((Vw, F, _color(rec["color"][i]))) | |
| return objs | |
| def interp_pose(R_k, t_k, times_k, t): | |
| """Pose at time t from key poses (linear translation, slerp rotation).""" | |
| from scipy.spatial.transform import Rotation, Slerp | |
| j = np.searchsorted(times_k, t, side="right") - 1 | |
| j = int(np.clip(j, 0, len(times_k) - 2)) if len(times_k) > 1 else 0 | |
| if len(times_k) == 1 or t <= times_k[0]: | |
| return R_k[0], t_k[0] | |
| if t >= times_k[-1]: | |
| return R_k[-1], t_k[-1] | |
| a = (t - times_k[j]) / (times_k[j + 1] - times_k[j]) | |
| rot = Slerp([0, 1], Rotation.from_matrix(np.stack([R_k[j], R_k[j + 1]])))([a]).as_matrix()[0] | |
| return rot, (1 - a) * t_k[j] + a * t_k[j + 1] | |
| # -------------------------------------------------------------------------------------- | |
| # Main entry point | |
| # -------------------------------------------------------------------------------------- | |
| def simulate(dataset: str = "MOVi-B", scene_source: str = "MOVi test scene", scene_index: int = 0, | |
| step_size: int = 1, num_frames: int = 480, num_objects: int = 6, seed: int = 0, | |
| progress=gr.Progress()): | |
| """Roll out RigidFormer on a multi-object rigid-body scene and render the prediction. | |
| Args: | |
| dataset: Which pretrained model / object set to use: "MOVi-B", "MOVi-A" or "MOVi-Sphere". | |
| scene_source: "MOVi test scene" (a held-out scene with simulator ground truth) or "Random new scene". | |
| scene_index: Index (0-119) of the held-out test scene, used when scene_source is "MOVi test scene". | |
| step_size: Integration step size in physical frames (1, 5 or 10; the simulation runs at 240 frames/s). | |
| num_frames: Rollout horizon in physical frames (480 = 2 seconds). | |
| num_objects: Number of objects for a random new scene (3-10). | |
| seed: Random seed for a random new scene. | |
| Returns: | |
| Rendered rollout video, per-frame error plot, a markdown summary and an .npz file with the predicted poses. | |
| """ | |
| step_size = int(step_size) | |
| if step_size not in (1, 5, 10): | |
| raise gr.Error("Step size must be 1, 5 or 10 (the step sizes RigidFormer was trained with).") | |
| num_frames = int(num_frames) | |
| num_frames = max(num_frames - num_frames % step_size, (WARMUP + 1) * step_size) | |
| if scene_source == "Random new scene": | |
| rec = random_scene(dataset, int(num_objects), int(seed), horizon=num_frames) | |
| has_gt = False | |
| else: | |
| recs = SCENES[dataset] | |
| rec = dict(recs[int(scene_index) % len(recs)]) | |
| rec["ground_truth"] = True | |
| has_gt = True | |
| n = rec["num_objects"] | |
| progress(0.05, desc="Building the point-cloud input") | |
| batch = build_batch(rec, step_size) | |
| num_steps = num_frames // step_size # network frames 0, s, ..., num_frames - s (as the evaluator) | |
| progress(0.15, desc=f"RigidFormer rollout ({num_steps - WARMUP} steps)") | |
| pred, elapsed = run_rollout(dataset, batch, num_steps) # (K, n, Nv, 3) | |
| K = pred.shape[0] | |
| times_k = np.arange(K) * step_size | |
| # Rigid poses of every object at every predicted frame (Kabsch fit to the predicted vertices) | |
| R_pred = np.zeros((K, n, 3, 3)); t_pred = np.zeros((K, n, 3)) | |
| for i in range(n): | |
| local = rec["size"][i] * COLLISION[rec["shape"][i]][0] | |
| nv = len(local) | |
| R_pred[:, i], t_pred[:, i] = kabsch_poses(local, pred[:, i, :nv]) | |
| # Errors vs. the simulator | |
| metrics_md = "" | |
| plot = None | |
| if has_gt: | |
| P_gt = rec["positions"][times_k] # (K, n, 3) | |
| R_gt = quat_wxyz_to_rotmat(rec["quaternions"][times_k]) | |
| pos_rmse = np.sqrt(((t_pred - P_gt) ** 2).sum(-1).mean(-1)) | |
| rel = np.einsum("tnij,tnlj->tnil", R_pred, R_gt) # R_pred @ R_gt^T | |
| cosang = (np.trace(rel, axis1=-2, axis2=-1) - 1) / 2 | |
| rot_err = np.degrees(np.sqrt((np.arccos(np.clip(cosang, -1, 1)) ** 2).mean(-1))) | |
| rows = [] | |
| for cp in (25, 50, 100, 200, 300, 400): | |
| if cp <= times_k[-1] and cp % step_size == 0 and cp // step_size >= WARMUP: | |
| k = cp // step_size | |
| rows.append(f"| @{cp} | {pos_rmse[k]:.4f} | {rot_err[k]:.2f} |") | |
| metrics_md = ("\n\n| physical frame | position RMSE (m) | rotation error (deg) |\n|---|---|---|\n" + "\n".join(rows) | |
| if rows else "") | |
| import matplotlib | |
| matplotlib.use("Agg") | |
| import matplotlib.pyplot as plt | |
| fig, ax1 = plt.subplots(figsize=(7, 3.2)) | |
| ax1.plot(times_k / FPS_PHYS, pos_rmse, color="#d97706", lw=2, label="position RMSE (m)") | |
| ax1.set_xlabel("time (s)"); ax1.set_ylabel("position RMSE (m)", color="#d97706") | |
| ax2 = ax1.twinx() | |
| ax2.plot(times_k / FPS_PHYS, rot_err, color="#475569", lw=2, ls="--", label="rotation error (deg)") | |
| ax2.set_ylabel("rotation error (deg)", color="#475569") | |
| ax1.set_title("RigidFormer vs. simulator ground truth") | |
| ax1.grid(alpha=0.3) | |
| fig.tight_layout() | |
| plot = fig | |
| # Video: ground truth (if any) | RigidFormer, one frame every 5 physical frames (real time at 48 fps) | |
| progress(0.4, desc="Rendering") | |
| out_path = os.path.join(tempfile.mkdtemp(), "rigidformer_rollout.mp4") | |
| writer = imageio.get_writer(out_path, fps=VIDEO_FPS, codec="libx264", quality=7, macro_block_size=16, | |
| ffmpeg_log_level="error") | |
| show_t = np.arange(0, times_k[-1] + 1, FPS_PHYS // VIDEO_FPS) | |
| R_q_all = quat_wxyz_to_rotmat(rec["quaternions"][: times_k[-1] + 1]) if has_gt else None | |
| for fi, t in enumerate(show_t): | |
| if fi % 8 == 0: | |
| progress(0.4 + 0.55 * fi / len(show_t), desc="Rendering") | |
| Rt = [interp_pose(R_pred[:, i], t_pred[:, i], times_k, t) for i in range(n)] | |
| R_now = np.stack([a for a, _ in Rt]); t_now = np.stack([b for _, b in Rt]) | |
| ctx = "context (ground truth)" if t <= times_k[WARMUP - 1] else "prediction" | |
| sub = f"t = {t / FPS_PHYS:.2f} s step size {step_size} [{ctx}]" | |
| right = render_frame(poses_to_objects(rec, R_now, t_now), "RigidFormer", sub) | |
| if has_gt: | |
| left = render_frame(poses_to_objects(rec, R_q_all[t], rec["positions"][t]), "Simulator (ground truth)", | |
| f"t = {t / FPS_PHYS:.2f} s") | |
| frame = np.concatenate([left, right], 1) | |
| else: | |
| frame = right | |
| writer.append_data(frame) | |
| writer.close() | |
| npz_path = os.path.join(tempfile.mkdtemp(), "rigidformer_poses.npz") | |
| np.savez(npz_path, rotation=R_pred.astype(np.float32), translation=t_pred.astype(np.float32), | |
| frame=times_k.astype(np.int32), shape=np.array(rec["shape"]), size=np.array(rec["size"], np.float32), | |
| material=np.array(rec["material"])) | |
| objs_desc = ", ".join(f"{m} {s}" for s, m in zip(rec["shape"], rec["material"])) | |
| summary = (f"**{dataset} — {rec['scene']}** · {n} objects ({objs_desc})\n\n" | |
| f"Rolled out {K - WARMUP} steps of size {step_size} ({times_k[-1] / FPS_PHYS:.2f} s of physics) " | |
| f"from {WARMUP} context frames in **{elapsed:.2f} s** on the GPU " | |
| f"({1000 * elapsed / max(K - WARMUP, 1):.1f} ms/step).") | |
| if not has_gt: | |
| summary += ("\n\nThis is a new scene with no simulator ground truth: only the first frames (free flight) are " | |
| "given, everything after that is predicted by RigidFormer.") | |
| summary += metrics_md | |
| return out_path, plot, summary, npz_path | |
| # -------------------------------------------------------------------------------------- | |
| # UI | |
| # -------------------------------------------------------------------------------------- | |
| DESCRIPTION = """ | |
| # RigidFormer: Learning Rigid Dynamics using Transformers | |
| **RigidFormer** ([paper](https://huggingface.co/papers/2605.09196) · [code](https://github.com/Frank-ZY-Dou/Dynamics-Modeling/tree/main/RigidFormer) · [weights](https://huggingface.co/frankzydou/RigidFormer)) is an object-centric Transformer that simulates multi-object rigid-body dynamics from **point sets** — no mesh connectivity. | |
| Given four context frames of each object's points, it predicts per-anchor accelerations, integrates them, and projects the update onto rigid motions (differentiable Kabsch), autoregressively for hundreds of steps at step size 1, 5 or 10. | |
| Pick a held-out **MOVi test scene** to compare against the physics simulator, or drop a **random new scene** of objects and watch the learned simulator take over. | |
| """ | |
| with gr.Blocks(title="RigidFormer") as demo: | |
| gr.Markdown(DESCRIPTION) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| dataset = gr.Dropdown(list(DATASETS), value="MOVi-B", label="Model / object set", | |
| info="MOVi-B: 11 KuBasic shapes · MOVi-A: cubes, cylinders, spheres · MOVi-Sphere: spheres") | |
| scene_source = gr.Radio(["MOVi test scene", "Random new scene"], value="MOVi test scene", label="Scene") | |
| scene_index = gr.Slider(0, 119, value=0, step=1, label="Test scene #", info="held-out scenes with simulator ground truth") | |
| with gr.Row(): | |
| num_objects = gr.Slider(3, 10, value=6, step=1, label="Objects (random scene)") | |
| seed = gr.Number(value=0, precision=0, label="Seed (random scene)") | |
| step_size = gr.Radio([1, 5, 10], value=1, label="Integration step size (physical frames)") | |
| with gr.Accordion("Advanced", open=False): | |
| num_frames = gr.Slider(60, 480, value=480, step=10, label="Horizon (physical frames @ 240 fps)") | |
| run = gr.Button("Simulate", variant="primary") | |
| with gr.Column(scale=2): | |
| video = gr.Video(label="Rollout", autoplay=True, loop=True) | |
| summary = gr.Markdown() | |
| plot = gr.Plot(label="Error vs. time") | |
| poses = gr.File(label="Predicted object poses (.npz)") | |
| inputs = [dataset, scene_source, scene_index, step_size, num_frames, num_objects, seed] | |
| outputs = [video, plot, summary, poses] | |
| run.click(simulate, inputs=inputs, outputs=outputs) | |
| gr.Examples( | |
| examples=[ | |
| ["MOVi-B", "MOVi test scene", 0, 1], | |
| ["MOVi-B", "MOVi test scene", 7, 5], | |
| ["MOVi-A", "MOVi test scene", 3, 1], | |
| ["MOVi-Sphere", "MOVi test scene", 11, 10], | |
| ["MOVi-B", "Random new scene", 0, 1], | |
| ], | |
| inputs=[dataset, scene_source, scene_index, step_size], | |
| outputs=outputs, fn=simulate, cache_examples=True, cache_mode="lazy", | |
| ) | |
| gr.Markdown("Scenes: [RigidFormer-MOVi](https://huggingface.co/datasets/frankzydou/RigidFormer-MOVi) test split " | |
| "(CC BY 4.0, derived from HopNet's MOVi datasets); meshes: Kubric KuBasic assets (CC BY-SA 4.0). " | |
| "Both are downloaded at startup.") | |
| if __name__ == "__main__": | |
| demo.launch(theme=gr.themes.Citrus(), mcp_server=True) | |