Emit meshable video pointmap grids

app.py

@@ -70,8 +70,7 @@ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 MAX_HEIGHT = 1024  # cap input height before processing — keeps everything fast
 VIDEO_MAX_HEIGHT = 512
 VIDEO_DEFAULT_FRAMES = 36
-VIDEO_DEFAULT_STRIDE = 5
-VIDEO_MAX_POINTS = 45_000
+VIDEO_DEFAULT_STRIDE = 3
 
 _fg_transform = transforms.Compose([
     transforms.Resize((1024, 768)),
@@ -292,19 +291,19 @@ def _sample_video_frames(video_path: str, max_frames: int) -> tuple[list[Image.I
     return frames, fps, source_w, source_h
 
 
-def _select_point_coords(mask: np.ndarray, stride: int, max_points: int) -> tuple[np.ndarray, np.ndarray]:
+def _sample_grid(mask: np.ndarray, stride: int) -> tuple[np.ndarray, np.ndarray, np.ndarray, int, int]:
     stride = int(max(1, min(stride, 16)))
-    grid = np.zeros_like(mask, dtype=bool)
-    grid[::stride, ::stride] = True
-    selected = mask & grid
-    yy, xx = np.where(selected)
-    if len(yy) < 128:
-        yy, xx = np.where(grid)
-    if len(yy) > max_points:
-        keep = np.linspace(0, len(yy) - 1, max_points, dtype=np.int64)
-        yy = yy[keep]
-        xx = xx[keep]
-    return yy.astype(np.int32), xx.astype(np.int32)
+    ys = np.arange(0, mask.shape[0], stride, dtype=np.int32)
+    xs = np.arange(0, mask.shape[1], stride, dtype=np.int32)
+    grid_y, grid_x = np.meshgrid(ys, xs, indexing="ij")
+    valid = mask[grid_y, grid_x].reshape(-1)
+    return (
+        grid_y.reshape(-1).astype(np.int32),
+        grid_x.reshape(-1).astype(np.int32),
+        valid.astype(np.uint8),
+        int(len(ys)),
+        int(len(xs)),
+    )
 
 
 def _point_sequence_zip(
@@ -322,6 +321,9 @@ def _point_sequence_zip(
     colors_frames: list[np.ndarray] = []
     sample_y: np.ndarray | None = None
     sample_x: np.ndarray | None = None
+    valid_mask: np.ndarray | None = None
+    grid_rows = 0
+    grid_cols = 0
     native_w = 0
     native_h = 0
 
@@ -332,14 +334,15 @@ def _point_sequence_zip(
         pointmap = _estimate_pointmap(image_bgr, model)
         native_h, native_w = pointmap.shape[:2]
 
-        if sample_y is None or sample_x is None:
+        if sample_y is None or sample_x is None or valid_mask is None:
             mask = _foreground_mask(image_pil, native_h, native_w)
-            sample_y, sample_x = _select_point_coords(mask, point_stride, VIDEO_MAX_POINTS)
-            print(f"[video] selected {len(sample_y)} points at native {native_w}x{native_h}")
+            sample_y, sample_x, valid_mask, grid_rows, grid_cols = _sample_grid(mask, point_stride)
+            print(f"[video] sampled {grid_cols}x{grid_rows} grid ({int(valid_mask.sum())} valid foreground points) at native {native_w}x{native_h}")
 
         image_native = np.array(image_pil.resize((native_w, native_h), Image.LANCZOS))
         points = pointmap[sample_y, sample_x].astype(np.float32)
         finite = np.isfinite(points).all(axis=1) & (points[:, 2] > 0.05) & (points[:, 2] < 25.0)
+        finite &= valid_mask.astype(bool)
         if finite.any():
             centroid = points[finite].mean(axis=0).astype(np.float32)
         else:
@@ -359,17 +362,19 @@ def _point_sequence_zip(
     positions = np.stack(positions_frames, axis=0).astype(np.float32)
     colors = np.stack(colors_frames, axis=0).astype(np.uint8)
     metadata = {
-        "format": "fpbox-sapiens-pointmap-sequence-v1",
+        "format": "fpbox-sapiens-pointmap-sequence-v2",
         "model": f"sapiens2-pointmap-{size}",
         "frameCount": int(positions.shape[0]),
         "fps": float(min(fps, max(1, positions.shape[0]))),
         "pointCount": int(positions.shape[1]),
+        "gridRows": int(grid_rows),
+        "gridCols": int(grid_cols),
         "width": int(native_w),
         "height": int(native_h),
         "sourceWidth": int(source_w),
         "sourceHeight": int(source_h),
-        "coordinateSystem": "x, -y, -z, centered per frame",
-        "dtype": {"positions": "float32", "colors": "uint8"},
+        "coordinateSystem": "x, -y, -z, centered per frame, regular sampled image grid",
+        "dtype": {"positions": "float32", "colors": "uint8", "valid": "uint8"},
     }
 
     out_path = tempfile.NamedTemporaryFile(delete=False, suffix=".zip").name
@@ -377,6 +382,7 @@ def _point_sequence_zip(
         zf.writestr("metadata.json", json.dumps(metadata, indent=2))
         zf.writestr("positions_f32.bin", positions.tobytes(order="C"))
         zf.writestr("colors_u8.bin", colors.tobytes(order="C"))
+        zf.writestr("valid_u8.bin", valid_mask.astype(np.uint8).tobytes(order="C"))
     return out_path