Compact flat-point GPU payload → frame cap 128 (denser); GLB built from points, verified == predictions_to_glb

app.py

@@ -37,7 +37,7 @@ CKPT_REPO = "robbyant/lingbot-map"
 CKPT_FILE = "lingbot-map-long.pt"
 NUM_SCALE_FRAMES = 8                   # initial bidirectional "scale" frames
 KEYFRAME_INTERVAL = 1                  # every frame is a keyframe — best quality for short clips
-MAX_FRAMES_HARD = 64                  # absolute cap (inference ~10s/48 frames → ample headroom under 120s)
+MAX_FRAMES_HARD = 128                 # absolute cap (~0.2s/frame → ~25s at 128, well under the 120s budget)
 GPU_DURATION = 120                    # seconds requested per GPU call (tune to your tier)
 
 # --------------------------------------------------------------------------------------
@@ -63,11 +63,10 @@ def gpu_reconstruct(images_cpu):
     if not _AGG_BF16_DONE["done"] and getattr(MODEL, "aggregator", None) is not None:
         MODEL.aggregator = MODEL.aggregator.to(dtype=torch.bfloat16)
         _AGG_BF16_DONE["done"] = True
-    return P.reconstruct_predictions(
+    return P.reconstruct_points(
         MODEL, images_cpu,
         num_scale_frames=NUM_SCALE_FRAMES,
         keyframe_interval=KEYFRAME_INTERVAL,
-        offload=True,
     )
 
 
@@ -108,26 +107,26 @@ def reconstruct(image_files, video_file, fps, max_frames, conf_thres, show_cam):
     images = P.preprocess_paths(paths)            # [S, 3, H, W] on CPU
     s, _, h, w = images.shape
 
-    vis = gpu_reconstruct(images)                 # GPU; returns NumPy vis-dict
+    rec = gpu_reconstruct(images)                 # GPU; returns a COMPACT flat point set
     t_infer = time.time() - t0
 
     t1 = time.time()
-    glb = P.export_glb(vis, conf_thres=conf_thres, show_cam=show_cam)
+    glb = P.build_glb_from_points(rec, conf_thres=conf_thres, show_cam=show_cam)
     t_export = time.time() - t1
 
-    kept, total = P.count_visible_points(vis, conf_thres)
+    kept, total = P.count_points(rec, conf_thres)
     status = _status_md(source, s, h, w, kept, total, t_infer, t_export, conf_thres)
-    return glb, vis, status
+    return glb, rec, status
 
 
-def reexport(vis, conf_thres, show_cam):
+def reexport(rec, conf_thres, show_cam):
     """CPU-only re-render when the user nudges the confidence filter / camera toggle.
-    Uses the cached vis-dict — no GPU, no re-inference."""
-    if not vis:
+    Uses the cached compact point set — no GPU, no re-inference."""
+    if not rec:
         return None, "Run a reconstruction first."
     t1 = time.time()
-    glb = P.export_glb(vis, conf_thres=conf_thres, show_cam=show_cam)
-    kept, total = P.count_visible_points(vis, conf_thres)
+    glb = P.build_glb_from_points(rec, conf_thres=conf_thres, show_cam=show_cam)
+    kept, total = P.count_points(rec, conf_thres)
     pct = (100.0 * kept / total) if total else 0.0
     msg = (
         f"Re-rendered (CPU). Points kept: **{kept:,} / {total:,}** ({pct:.0f}%) at "
@@ -163,8 +162,8 @@ def build_demo():
                 with gr.Accordion("Settings", open=True):
                     fps = gr.Slider(1, 12, value=6, step=1,
                                     label="Video sampling FPS (frames/sec to extract)")
-                    max_frames = gr.Slider(2, MAX_FRAMES_HARD, value=48, step=1,
-                                           label="Max frames (spread evenly across the whole video)")
+                    max_frames = gr.Slider(2, MAX_FRAMES_HARD, value=80, step=1,
+                                           label="Max frames — higher = denser & more complete (spread across the whole video)")
                     conf_thres = gr.Slider(0, 95, value=60, step=1,
                                            label="Confidence filter — drops the lowest-confidence % of points (raise to declutter)")
                     show_cam = gr.Checkbox(value=True, label="Show camera trajectory")

pipeline.py

@@ -360,39 +360,86 @@ def reconstruct_predictions(
 # GLB export (CPU only — no GPU needed)
 # =============================================================================
 
-def _add_trajectory_line(scene, extrinsic, seed: int = 0):
-    """Add a single clean colored tube tracing the camera path (README "track" style),
-    in the SAME aligned frame predictions_to_glb uses. Far less cluttered than drawing a
-    frustum per frame. Colors run start→end along gist_rainbow."""
-    import trimesh
-    from scipy.spatial.transform import Rotation
-    import matplotlib
-
-    extrinsic = np.asarray(extrinsic, dtype=np.float64)
-    S = len(extrinsic)
-    if S < 2:
-        return scene
+def _camera_centers_world(extrinsic):
+    """(camera centers in world coords [S,3], w2c 4x4 [S,4,4]) from w2c [S,3,4]."""
+    extr = np.asarray(extrinsic, dtype=np.float64)
+    S = len(extr)
     w2c = np.repeat(np.eye(4)[None], S, axis=0)
-    w2c[:, :3, :4] = extrinsic
-    centers = np.linalg.inv(w2c)[:, :3, 3]                       # world camera centers (c2w t)
-    # Replicate predictions_to_glb.apply_scene_alignment so the line overlays the cloud.
+    w2c[:, :3, :4] = extr
+    centers = np.linalg.inv(w2c)[:, :3, 3]
+    return centers, w2c
+
+
+def _alignment_transform(w2c):
+    """The same scene alignment predictions_to_glb applies (orient to the first camera +
+    flip into the GLB/OpenGL convention). Applied to the WHOLE scene (points + trajectory).
+    Verified to match predictions_to_glb's apply_scene_alignment."""
+    from scipy.spatial.transform import Rotation
     opengl = np.eye(4); opengl[1, 1] = -1.0; opengl[2, 2] = -1.0
     align = np.eye(4); align[:3, :3] = Rotation.from_euler("y", 180, degrees=True).as_matrix()
-    initial = np.linalg.inv(w2c[0]) @ opengl @ align
-    pts = (np.c_[centers, np.ones(S)] @ initial.T)[:, :3]        # aligned camera centers
-    ext = float(np.linalg.norm(pts.max(0) - pts.min(0))) or 1.0
+    return np.linalg.inv(w2c[0]) @ opengl @ align
+
+
+def _add_trajectory_tube(scene, centers_world):
+    """Add a clean colored tube tracing the camera path (README "track" style) in WORLD
+    coords; the caller applies the scene alignment afterwards. gist_rainbow start→end."""
+    import trimesh
+    import matplotlib
+    S = len(centers_world)
+    if S < 2:
+        return
+    ext = float(np.linalg.norm(centers_world.max(0) - centers_world.min(0))) or 1.0
     radius = max(ext * 0.004, 1e-5)
     cmap = matplotlib.colormaps.get_cmap("gist_rainbow")
     tubes = []
     for i in range(S - 1):
-        if np.linalg.norm(pts[i + 1] - pts[i]) < 1e-9:
+        if np.linalg.norm(centers_world[i + 1] - centers_world[i]) < 1e-9:
             continue
-        seg = trimesh.creation.cylinder(radius=radius, segment=pts[i:i + 2], sections=8)
+        seg = trimesh.creation.cylinder(radius=radius, segment=centers_world[i:i + 2], sections=8)
         seg.visual.vertex_colors = tuple(int(255 * x) for x in cmap(i / (S - 1))[:3]) + (255,)
         tubes.append(seg)
     if tubes:
         scene.add_geometry(trimesh.util.concatenate(tubes))
-    return scene
+
+
+def _flatten_vis(vis_predictions: dict) -> dict:
+    """Flatten a vis dict → a compact flat point set in WORLD coords (the format the GPU
+    call returns and the GLB builder consumes). Unprojects depth at full resolution.
+    Returns {points (N,3) f32, colors (N,3) u8, conf (N,) f32, extrinsic (S,3,4) f32}.
+    """
+    vis = vis_predictions
+    if "world_points" in vis:
+        world = np.asarray(vis["world_points"], np.float32)
+        conf = np.asarray(vis.get("world_points_conf"), np.float32)
+    else:
+        world = np.asarray(
+            unproject_depth_map_to_point_map(vis["depth"], vis["extrinsic"], vis["intrinsic"]),
+            np.float32,
+        )
+        conf = np.asarray(vis["depth_conf"], np.float32)
+    imgs = np.asarray(vis["images"])
+    if imgs.ndim == 4 and imgs.shape[1] == 3:          # NCHW → NHWC
+        imgs = np.transpose(imgs, (0, 2, 3, 1))
+    return {
+        "points": world.reshape(-1, 3),
+        "colors": np.clip(imgs.reshape(-1, 3) * 255.0, 0, 255).astype(np.uint8),
+        "conf": conf.reshape(-1),
+        "extrinsic": np.asarray(vis["extrinsic"], np.float32),
+    }
+
+
+def _downsample_rec(rec: dict, max_points: int, seed: int = 0) -> dict:
+    """Drop invalid points and randomly cap to max_points so the GPU→app payload stays
+    bounded and independent of frame count."""
+    pts, cols, conf = rec["points"], rec["colors"], rec["conf"]
+    valid = conf > 1e-5
+    pts, cols, conf = pts[valid], cols[valid], conf[valid]
+    n_total = int(pts.shape[0])
+    if n_total > max_points:
+        sel = np.random.default_rng(seed).choice(n_total, max_points, replace=False)
+        pts, cols, conf = pts[sel], cols[sel], conf[sel]
+    return {"points": pts, "colors": cols, "conf": conf,
+            "extrinsic": rec["extrinsic"], "n_total": n_total}
 
 
 def export_glb(
@@ -403,66 +450,89 @@ def export_glb(
     mask_sky: bool = False,
     target_dir: Optional[str] = None,
 ) -> str:
-    """Export the vis dict to a .glb and return its path.
+    """Convenience wrapper (smoke_test / local CPU runs): flatten a vis dict → build the GLB.
+    The Space uses reconstruct_points() + build_glb_from_points() directly (compact payload)."""
+    rec = _flatten_vis(vis_predictions)
+    return build_glb_from_points(rec, out_path=out_path, conf_thres=conf_thres, show_cam=show_cam)
 
-    The streaming model (GCTStream) is built with ``enable_point=False``, so it emits
-    **depth**, not a point map. We unproject depth → world points (``world_points_from_depth``,
-    via the camera w2c extrinsics + intrinsics) and use the exporter's *Depthmap* branch,
-    which filters by ``depth_conf``. If a point map is present (e.g. a future enable_point
-    build), we use it directly via the *Pointmap* branch instead.
 
-    Reminder: ``conf_thres`` is a *percentile* — points below the ``conf_thres``-th
-    percentile of confidence are dropped. 0 keeps everything; higher = fewer points.
+def build_glb_from_points(
+    rec: dict,
+    out_path: Optional[str] = None,
+    conf_thres: float = 60.0,
+    show_cam: bool = True,
+    display_max: int = 2_000_000,
+    seed: int = 0,
+) -> str:
+    """Build the .glb from a flat point set: percentile-filter by confidence, downsample for
+    smooth in-browser rendering, add the camera-path line, and apply the scene alignment.
+
+    ``conf_thres`` is a *percentile* — points below the conf_thres-th percentile of
+    confidence are dropped. 0 keeps everything; higher = fewer points.
     """
+    import trimesh
     if out_path is None:
         out_path = os.path.join(tempfile.mkdtemp(prefix="lingbot_glb_"), "scene.glb")
 
-    preds = dict(vis_predictions)  # shallow copy; we may add world_points_from_depth
-    if "world_points" in preds:
-        prediction_mode = "Predicted Pointmap"
-    else:
-        preds["world_points_from_depth"] = unproject_depth_map_to_point_map(
-            preds["depth"], preds["extrinsic"], preds["intrinsic"]
-        )
-        prediction_mode = "Predicted Depthmap"
-
-    scene = predictions_to_glb(
-        preds,
-        conf_thres=conf_thres,
-        show_cam=False,            # draw a single clean trajectory line instead of per-frame frustums
-        mask_sky=mask_sky,
-        target_dir=target_dir,
-        prediction_mode=prediction_mode,
-    )
-    if show_cam and "extrinsic" in preds:
-        _add_trajectory_line(scene, preds["extrinsic"])
-
-    # Downsample the point cloud so the browser viewer (model-viewer) stays smooth — it
-    # bogs down past a few million points. Camera frustums are left intact.
-    import trimesh
-    display_max = 2_000_000
-    for _name in list(scene.geometry.keys()):
-        _g = scene.geometry[_name]
-        if isinstance(_g, trimesh.PointCloud) and len(_g.vertices) > display_max:
-            _sel = np.random.default_rng(0).choice(len(_g.vertices), display_max, replace=False)
-            _cols = np.asarray(_g.colors) if _g.colors is not None else None
-            _cols = _cols[_sel] if (_cols is not None and len(_cols) == len(_g.vertices)) else None
-            scene.geometry[_name] = trimesh.PointCloud(vertices=np.asarray(_g.vertices)[_sel], colors=_cols)
-
+    pts = np.asarray(rec["points"], np.float64)
+    cols = np.asarray(rec["colors"])
+    conf = np.asarray(rec["conf"], np.float64)
+
+    if conf.size:
+        thr = np.percentile(conf, conf_thres) if conf_thres > 0 else 0.0
+        mask = (conf >= thr) & (conf > 1e-5)
+        pts, cols = pts[mask], cols[mask]
+    if len(pts) > display_max:                          # browser viewer bogs down past a few M points
+        sel = np.random.default_rng(seed).choice(len(pts), display_max, replace=False)
+        pts, cols = pts[sel], cols[sel]
+    if len(pts) == 0:                                   # degenerate guard
+        pts, cols = np.zeros((1, 3)), np.array([[255, 255, 255]])
+
+    scene = trimesh.Scene()
+    scene.add_geometry(trimesh.PointCloud(vertices=pts, colors=cols))
+    centers, w2c = _camera_centers_world(rec["extrinsic"])
+    if show_cam:
+        _add_trajectory_tube(scene, centers)
+    scene.apply_transform(_alignment_transform(w2c))    # align points + trajectory together
     scene.export(out_path)
     return out_path
 
 
-def count_visible_points(vis_predictions: dict, conf_thres: float = 50.0) -> Tuple[int, int]:
-    """(kept, total) point counts after the same percentile filter the exporter uses.
-    Uses world_points_conf in Pointmap mode, else depth_conf (the streaming/depth path).
-    Handy for the status readout so users can see when conf_thres nukes the cloud.
-    """
-    conf_key = "world_points_conf" if "world_points" in vis_predictions else "depth_conf"
-    conf = np.asarray(vis_predictions[conf_key]).reshape(-1)
+def count_points(rec: dict, conf_thres: float = 60.0) -> Tuple[int, int]:
+    """(kept, total) after the same percentile filter build_glb_from_points uses."""
+    conf = np.asarray(rec["conf"]).reshape(-1)
     total = conf.size
     if total == 0:
         return 0, 0
     thr = np.percentile(conf, conf_thres) if conf_thres > 0 else 0.0
     kept = int(np.count_nonzero((conf >= thr) & (conf > 1e-5)))
     return kept, total
+
+
+def count_visible_points(vis_predictions: dict, conf_thres: float = 60.0) -> Tuple[int, int]:
+    """Back-compat wrapper for the vis-dict form (smoke_test)."""
+    return count_points(_flatten_vis(vis_predictions), conf_thres)
+
+
+# =============================================================================
+# GPU entry point that returns a COMPACT, frame-count-independent payload
+# =============================================================================
+
+def reconstruct_points(
+    model: GCTStream,
+    images: torch.Tensor,
+    num_scale_frames: int = 8,
+    keyframe_interval: int = 1,
+    max_return_points: int = 6_000_000,
+) -> dict:
+    """End-to-end GPU path → a COMPACT flat point set whose size is bounded (independent of
+    frame count). This is what lets the Space return many frames across the @spaces.GPU
+    boundary without a huge per-pixel payload.
+
+    Returns {points (M,3) f32, colors (M,3) u8, conf (M,) f32, extrinsic (S,3,4) f32, n_total}.
+    """
+    vis = reconstruct_predictions(
+        model, images, num_scale_frames=num_scale_frames,
+        keyframe_interval=keyframe_interval, offload=True,
+    )
+    return _downsample_rec(_flatten_vis(vis), max_return_points)