feat: video walk-around -> 3D environment map (SIFT + metric depth scale recovery)

app.py

@@ -2,6 +2,7 @@ import spaces  # MUST be the very first import for ZeroGPU to work
 
 import io
 import numpy as np
+import cv2
 from PIL import Image, UnidentifiedImageError
 import torch
 import gradio as gr
@@ -28,7 +29,7 @@ print(f"Model successfully loaded on target device: {device}")
 
 
 # =====================================================================
-# 2. Core inference helper (shared by Gradio UI and the REST endpoint)
+# 2. Shared helpers
 # =====================================================================
 def _run_depth_inference(image: Image.Image):
     """Returns post-processed depth output list. Caller handles indexing."""
@@ -41,55 +42,8 @@ def _run_depth_inference(image: Image.Image):
     )
 
 
-# =====================================================================
-# 3. Gradio function - @spaces.GPU tells ZeroGPU to allocate a GPU here.
-#    This MUST be bound to a Gradio event handler for ZeroGPU to detect it.
-# =====================================================================
-@spaces.GPU(duration=120)
-def create_3d_point_cloud(input_pil_image):
-    """Processes the image, calculates 3D vertices, and writes a binary .ply file."""
-    if input_pil_image is None:
-        return None
-
-    # Step A: Downscale large images before inference.
-    # DepthPro quality is stable at 512px; this dramatically cuts runtime.
-    image = input_pil_image.convert("RGB")
-    MAX_DIM = 512
-    if max(image.width, image.height) > MAX_DIM:
-        image.thumbnail((MAX_DIM, MAX_DIM), Image.LANCZOS)
-
-    post_processed = _run_depth_inference(image)
-    depth_matrix = (
-        post_processed[0]["predicted_depth"]
-        .detach().cpu().numpy().astype(np.float32)
-    )
-
-    # Step B: Pinhole camera projection (u, v, depth) -> (X, Y, Z)
-    height, width = depth_matrix.shape
-    color_np = np.array(image, dtype=np.float32) / 255.0
-
-    focal_length = float(max(width, height))
-    cx, cy = width / 2.0, height / 2.0
-    u, v = np.meshgrid(np.arange(width, dtype=np.float32),
-                        np.arange(height, dtype=np.float32))
-
-    z = depth_matrix.flatten()
-    x = ((u.flatten() - cx) * z) / focal_length
-    y = ((v.flatten() - cy) * z) / focal_length
-
-    # Keep objects between 10 cm and 30 m; discard sky / infinite background
-    valid_mask = (z > 0.1) & (z < 30.0)
-    x, y, z = x[valid_mask], y[valid_mask], z[valid_mask]
-    rgb_ints = (color_np.reshape(-1, 3)[valid_mask] * 255).astype(np.uint8)
-
-    # Step C: Subsample to keep the browser renderer fast (<= 200k points)
-    MAX_POINTS = 200_000
-    if len(z) > MAX_POINTS:
-        idx = np.random.choice(len(z), MAX_POINTS, replace=False)
-        x, y, z = x[idx], y[idx], z[idx]
-        rgb_ints = rgb_ints[idx]
-
-    # Step D: Write BINARY PLY via numpy (100x faster than an ASCII for-loop)
+def _write_binary_ply(x, y, z, rgb_ints, path: str) -> str:
+    """Pack arrays into a binary little-endian PLY and write to path."""
     num_points = len(z)
     vertex_dtype = np.dtype([
         ('x', '<f4'), ('y', '<f4'), ('z', '<f4'),
@@ -97,13 +51,13 @@ def create_3d_point_cloud(input_pil_image):
     ])
     vertices = np.empty(num_points, dtype=vertex_dtype)
     vertices['x']     = x.astype(np.float32)
-    vertices['y']     = (-y).astype(np.float32)   # invert Y so it loads right-side up
+    vertices['y']     = (-y).astype(np.float32)   # invert Y → right-side up
     vertices['z']     = z.astype(np.float32)
     vertices['red']   = rgb_ints[:, 0]
     vertices['green'] = rgb_ints[:, 1]
     vertices['blue']  = rgb_ints[:, 2]
 
-    ply_header = (
+    header = (
         "ply\n"
         "format binary_little_endian 1.0\n"
         f"element vertex {num_points}\n"
@@ -115,42 +69,306 @@ def create_3d_point_cloud(input_pil_image):
         "property uchar blue\n"
         "end_header\n"
     )
-
-    ply_filename = "/tmp/output_scene.ply"
-    with open(ply_filename, "wb") as f:
-        f.write(ply_header.encode("ascii"))
+    with open(path, "wb") as f:
+        f.write(header.encode("ascii"))
         f.write(vertices.tobytes())
+    return path
+
+
+def _frame_to_pointcloud(image: Image.Image, depth: np.ndarray, max_pts: int = 50_000):
+    """Project a single frame into (x, y, z, rgb) arrays in camera space."""
+    h, w = depth.shape
+    color_np = np.array(image, dtype=np.float32) / 255.0
+    focal = float(max(w, h))
+    cx, cy = w / 2.0, h / 2.0
+
+    u_g, v_g = np.meshgrid(np.arange(w, dtype=np.float32),
+                             np.arange(h, dtype=np.float32))
+    z = depth.flatten()
+    x = ((u_g.flatten() - cx) * z) / focal
+    y = ((v_g.flatten() - cy) * z) / focal
+
+    valid = (z > 0.1) & (z < 30.0)
+    x, y, z = x[valid], y[valid], z[valid]
+    rgb = (color_np.reshape(-1, 3)[valid] * 255).astype(np.uint8)
 
-    return ply_filename
+    if len(z) > max_pts:
+        idx = np.random.choice(len(z), max_pts, replace=False)
+        x, y, z, rgb = x[idx], y[idx], z[idx], rgb[idx]
+
+    return x, y, z, rgb
 
 
 # =====================================================================
-# 4. Gradio UI - uses demo.launch() (NOT gr.mount_gradio_app) so that
-#    ZeroGPU can properly detect the @spaces.GPU decorated function above.
+# 3a. Single-image function (original feature)
 # =====================================================================
-with gr.Blocks(title="Software-ToF 3D Cloud Visualizer") as demo:
-    gr.Markdown("# 📷 Software-ToF Camera Simulation")
-    gr.Markdown(
-        "Upload any image to calculate absolute metric depths and generate "
-        "an interactive 3D Point Cloud visualization."
+@spaces.GPU(duration=120)
+def create_3d_point_cloud(input_pil_image):
+    """Single photo → coloured 3D point cloud PLY."""
+    if input_pil_image is None:
+        return None
+
+    image = input_pil_image.convert("RGB")
+    if max(image.width, image.height) > 512:
+        image.thumbnail((512, 512), Image.LANCZOS)
+
+    post = _run_depth_inference(image)
+    depth = post[0]["predicted_depth"].detach().cpu().numpy().astype(np.float32)
+
+    x, y, z, rgb = _frame_to_pointcloud(image, depth, max_pts=200_000)
+    return _write_binary_ply(x, y, z, rgb, "/tmp/output_scene.ply")
+
+
+# =====================================================================
+# 3b. Video → environment map (new feature)
+# =====================================================================
+@spaces.GPU(duration=300)
+def create_map_from_video(video_path, frame_stride: int = 10, max_frames: int = 20):
+    """
+    Walk-around video → merged 3D environment map.
+
+    Pipeline
+    --------
+    1. Extract keyframes (every `frame_stride` video frames, up to `max_frames`).
+    2. Downscale each frame to 512 px and run DepthPro → metric depth.
+    3. SIFT feature detection + FLANN matching between consecutive frames.
+    4. Essential Matrix → R, t̂  (unit translation from recoverPose).
+    5. Use metric depth at matched inliers to recover the true translation scale.
+    6. Accumulate 4×4 camera poses in a world coordinate frame.
+    7. Project each frame's pixels into world space and merge point clouds.
+    8. Subsample to ≤ 500k points and write binary PLY.
+    """
+    if video_path is None:
+        return None, "No video provided."
+
+    # ------------------------------------------------------------------
+    # Step 1: Extract keyframes
+    # ------------------------------------------------------------------
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        return None, "Could not open video file."
+
+    frames_pil = []
+    total_video_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    frame_indices = range(0, total_video_frames, int(frame_stride))
+
+    for fi in frame_indices:
+        if len(frames_pil) >= int(max_frames):
+            break
+        cap.set(cv2.CAP_PROP_POS_FRAMES, fi)
+        ret, bgr = cap.read()
+        if not ret:
+            break
+        rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
+        pil = Image.fromarray(rgb)
+        if max(pil.width, pil.height) > 512:
+            pil.thumbnail((512, 512), Image.LANCZOS)
+        frames_pil.append(pil)
+    cap.release()
+
+    n = len(frames_pil)
+    if n < 2:
+        return None, f"Need at least 2 frames; got {n}."
+
+    W, H = frames_pil[0].width, frames_pil[0].height
+    focal = float(max(W, H))
+    cx, cy = W / 2.0, H / 2.0
+    K = np.array([[focal, 0, cx], [0, focal, cy], [0, 0, 1]], dtype=np.float64)
+
+    # ------------------------------------------------------------------
+    # Step 2: DepthPro inference on every keyframe
+    # ------------------------------------------------------------------
+    depth_maps = []
+    for pil in frames_pil:
+        post = _run_depth_inference(pil)
+        d = post[0]["predicted_depth"].detach().cpu().numpy().astype(np.float32)
+        depth_maps.append(d)
+
+    # ------------------------------------------------------------------
+    # Step 3: SIFT keypoints for every frame
+    # ------------------------------------------------------------------
+    sift = cv2.SIFT_create(nfeatures=2000)
+    kps_list, desc_list = [], []
+    for pil in frames_pil:
+        gray = cv2.cvtColor(np.array(pil), cv2.COLOR_RGB2GRAY)
+        kps, desc = sift.detectAndCompute(gray, None)
+        kps_list.append(kps)
+        desc_list.append(desc)
+
+    # ------------------------------------------------------------------
+    # Step 4-5: Relative pose estimation with metric-depth scale recovery
+    # ------------------------------------------------------------------
+    flann = cv2.FlannBasedMatcher({"algorithm": 1, "trees": 5}, {"checks": 50})
+
+    # poses[i] = 4×4 world-to-camera transform for frame i
+    poses = [np.eye(4, dtype=np.float64)]
+
+    for i in range(n - 1):
+        d_i   = desc_list[i]
+        d_ip1 = desc_list[i + 1]
+
+        # Fall back to identity step if SIFT failed on either frame
+        if d_i is None or d_ip1 is None or len(d_i) < 10 or len(d_ip1) < 10:
+            poses.append(poses[-1].copy())
+            continue
+
+        try:
+            matches = flann.knnMatch(d_i, d_ip1, k=2)
+        except cv2.error:
+            poses.append(poses[-1].copy())
+            continue
+
+        # Lowe's ratio test
+        good = [m for m, nn in matches if m.distance < 0.7 * nn.distance]
+        if len(good) < 10:
+            poses.append(poses[-1].copy())
+            continue
+
+        pts1 = np.float32([kps_list[i][m.queryIdx].pt for m in good])
+        pts2 = np.float32([kps_list[i + 1][m.trainIdx].pt for m in good])
+
+        E, mask_E = cv2.findEssentialMat(
+            pts1, pts2, K, method=cv2.RANSAC, prob=0.999, threshold=1.0
+        )
+        if E is None:
+            poses.append(poses[-1].copy())
+            continue
+
+        _, R, t_unit, mask_pose = cv2.recoverPose(E, pts1, pts2, K, mask=mask_E)
+
+        # --- Scale recovery via metric depth ---
+        # recoverPose gives ||t|| = 1. We use DepthPro depth at inlier
+        # correspondences to estimate the true translation magnitude.
+        inliers = mask_pose.ravel() == 255
+        pts1_in = pts1[inliers]
+        pts2_in = pts2[inliers]
+
+        depth_i   = depth_maps[i]
+        depth_ip1 = depth_maps[i + 1]
+        h_d, w_d  = depth_i.shape
+
+        scales = []
+        for p1, p2 in zip(pts1_in, pts2_in):
+            u1, v1 = int(np.clip(p1[0], 0, w_d - 1)), int(np.clip(p1[1], 0, h_d - 1))
+            u2, v2 = int(np.clip(p2[0], 0, w_d - 1)), int(np.clip(p2[1], 0, h_d - 1))
+
+            d1 = float(depth_i[v1, u1])
+            d2 = float(depth_ip1[v2, u2])
+            if not (0.1 < d1 < 30.0 and 0.1 < d2 < 30.0):
+                continue
+
+            # 3-D positions in camera-i and camera-(i+1) coords
+            P1 = np.array([(p1[0] - cx) * d1 / focal,
+                            (p1[1] - cy) * d1 / focal,
+                            d1])
+            P2 = np.array([(p2[0] - cx) * d2 / focal,
+                            (p2[1] - cy) * d2 / focal,
+                            d2])
+            # Actual translation ≈ P2 - R @ P1
+            t_actual = P2 - R @ P1
+            scales.append(np.linalg.norm(t_actual))
+
+        scale = float(np.median(scales)) if scales else 1.0
+        t_scaled = t_unit * scale
+
+        # Build incremental 4×4 transform
+        T_rel = np.eye(4, dtype=np.float64)
+        T_rel[:3, :3] = R
+        T_rel[:3, 3]  = t_scaled.ravel()
+
+        poses.append(T_rel @ poses[-1])
+
+    # ------------------------------------------------------------------
+    # Step 6-7: Project each frame → world space and accumulate
+    # ------------------------------------------------------------------
+    all_pts, all_rgb = [], []
+
+    for i, (pil, depth, pose) in enumerate(zip(frames_pil, depth_maps, poses)):
+        x_c, y_c, z_c, rgb = _frame_to_pointcloud(pil, depth, max_pts=50_000)
+
+        # Camera → world: pts_world = inv(pose) @ pts_cam_homogeneous
+        pts_h = np.vstack([x_c, y_c, z_c, np.ones_like(z_c)])   # 4 × N
+        pts_w = (np.linalg.inv(pose) @ pts_h)[:3]                # 3 × N
+
+        all_pts.append(pts_w.T)   # N × 3
+        all_rgb.append(rgb)
+
+    XYZ = np.concatenate(all_pts, axis=0)
+    RGB = np.concatenate(all_rgb, axis=0)
+
+    # ------------------------------------------------------------------
+    # Step 8: Global subsample → binary PLY
+    # ------------------------------------------------------------------
+    MAX_TOTAL = 500_000
+    if len(XYZ) > MAX_TOTAL:
+        idx = np.random.choice(len(XYZ), MAX_TOTAL, replace=False)
+        XYZ = XYZ[idx]
+        RGB = RGB[idx]
+
+    ply_path = _write_binary_ply(XYZ[:, 0], XYZ[:, 1], XYZ[:, 2],
+                                  RGB, "/tmp/environment_map.ply")
+    msg = (
+        f"✅ Processed {n} frames · "
+        f"{len(XYZ):,} points in world map"
     )
+    return ply_path, msg
 
-    with gr.Row():
-        with gr.Column():
-            input_img  = gr.Image(type="pil", label="Input Scene Photo")
-            submit_btn = gr.Button("Generate ToF 3D Structure", variant="primary")
-        with gr.Column():
-            output_3d = gr.Model3D(
-                label="Interactive 3D Point Cloud Canvas",
-                display_mode="point_cloud",
-            )
 
-    submit_btn.click(fn=create_3d_point_cloud, inputs=[input_img], outputs=[output_3d])
+# =====================================================================
+# 4. Gradio UI — two tabs
+# =====================================================================
+with gr.Blocks(title="Software-ToF 3D Cloud Visualizer") as demo:
+    gr.Markdown("# 📷 Software-ToF Camera Simulation")
+
+    with gr.Tab("📷 Single Image → 3D Snapshot"):
+        gr.Markdown(
+            "Upload any photo to get an interactive 3D point cloud "
+            "with absolute metric depths from Apple DepthPro."
+        )
+        with gr.Row():
+            with gr.Column():
+                input_img  = gr.Image(type="pil", label="Input Scene Photo")
+                img_btn    = gr.Button("Generate ToF 3D Structure", variant="primary")
+            with gr.Column():
+                output_3d  = gr.Model3D(
+                    label="Interactive 3D Point Cloud",
+                    display_mode="point_cloud",
+                )
+        img_btn.click(fn=create_3d_point_cloud,
+                      inputs=[input_img],
+                      outputs=[output_3d])
+
+    with gr.Tab("🎥 Video → Environment Map"):
+        gr.Markdown(
+            "**Walk around a room while recording, then upload the video.**  \n"
+            "Frames are matched with SIFT feature correspondences. "
+            "DepthPro provides metric depth per frame; the scale of each "
+            "relative camera translation is recovered from matched 3-D points, "
+            "then all frames are fused into a single world-coordinate map."
+        )
+        with gr.Row():
+            with gr.Column():
+                video_input   = gr.Video(label="Walk-around Video")
+                frame_stride  = gr.Slider(5, 30, value=10, step=5,
+                                           label="Frame Stride (sample every N frames)")
+                max_frames    = gr.Slider(5, 30, value=20, step=5,
+                                           label="Max Frames to Process")
+                video_btn     = gr.Button("Generate Environment Map", variant="primary")
+            with gr.Column():
+                map_output    = gr.Model3D(
+                    label="3D Environment Map",
+                    display_mode="point_cloud",
+                )
+                status_box    = gr.Textbox(label="Status", interactive=False)
+        video_btn.click(
+            fn=create_map_from_video,
+            inputs=[video_input, frame_stride, max_frames],
+            outputs=[map_output, status_box],
+        )
 
 
 # =====================================================================
-# 5. Add the REST /get_metric_depth endpoint via Gradio's internal
-#    FastAPI app - this avoids gr.mount_gradio_app which breaks ZeroGPU.
+# 5. REST endpoints via Gradio's internal FastAPI app
 # =====================================================================
 @demo.app.get("/status")
 def status():
@@ -161,31 +379,22 @@ def status():
 async def get_metric_depth(file: UploadFile = File(...)):
     if not (file.content_type or "").startswith("image/"):
         raise HTTPException(status_code=400, detail="Uploaded file must be an image.")
-
     try:
-        image_bytes = await file.read()
-        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        image = Image.open(io.BytesIO(await file.read())).convert("RGB")
     except (UnidentifiedImageError, OSError):
         raise HTTPException(status_code=400, detail="Could not decode image.")
-
     try:
-        post_processed = _run_depth_inference(image)
-        metric_depth_map = (
-            post_processed[0]["predicted_depth"]
-            .detach().cpu().numpy().astype(np.float32)
-        )
-        buffer = io.BytesIO()
-        np.save(buffer, metric_depth_map)
-        return Response(content=buffer.getvalue(), media_type="application/octet-stream")
-
+        post = _run_depth_inference(image)
+        depth = post[0]["predicted_depth"].detach().cpu().numpy().astype(np.float32)
+        buf = io.BytesIO()
+        np.save(buf, depth)
+        return Response(content=buf.getvalue(), media_type="application/octet-stream")
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Inference error: {str(e)}")
 
 
 # =====================================================================
-# 6. Launch - Hugging Face auto-detects the `demo` variable when
-#    sdk: gradio is set and launches it. The explicit launch() call
-#    here ensures it also works when run locally.
+# 6. Launch
 # =====================================================================
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()

requirements.txt

@@ -8,3 +8,4 @@ gradio
 fastapi
 uvicorn[standard]
 python-multipart
+opencv-python-headless