Spaces:
Sleeping
Sleeping
| 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 | |
| from fastapi import File, UploadFile, HTTPException | |
| from fastapi.responses import Response | |
| from transformers import DepthProImageProcessor, DepthProForDepthEstimation | |
| # ===================================================================== | |
| # 1. Load Apple Depth Pro at module level. | |
| # ZeroGPU safely emulates CUDA during startup so .to("cuda") works here. | |
| # ===================================================================== | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| dtype = torch.float16 if torch.cuda.is_available() else torch.float32 | |
| print("Loading Apple Depth Pro components...") | |
| image_processor = DepthProImageProcessor.from_pretrained("apple/DepthPro-hf") | |
| model = DepthProForDepthEstimation.from_pretrained( | |
| "apple/DepthPro-hf", | |
| torch_dtype=dtype, | |
| ).to(device) | |
| model.eval() | |
| print(f"Model successfully loaded on target device: {device}") | |
| # ===================================================================== | |
| # 2. Shared helpers | |
| # ===================================================================== | |
| def _run_depth_inference(image: Image.Image): | |
| """Returns post-processed depth output list. Caller handles indexing.""" | |
| inputs = image_processor(images=image, return_tensors="pt").to(device) | |
| with torch.no_grad(): | |
| outputs = model(**inputs) | |
| return image_processor.post_process_depth_estimation( | |
| outputs, | |
| target_sizes=[(image.height, image.width)], | |
| ) | |
| 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'), | |
| ('red', 'u1'), ('green', 'u1'), ('blue', 'u1'), | |
| ]) | |
| vertices = np.empty(num_points, dtype=vertex_dtype) | |
| vertices['x'] = x.astype(np.float32) | |
| 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] | |
| header = ( | |
| "ply\n" | |
| "format binary_little_endian 1.0\n" | |
| f"element vertex {num_points}\n" | |
| "property float x\n" | |
| "property float y\n" | |
| "property float z\n" | |
| "property uchar red\n" | |
| "property uchar green\n" | |
| "property uchar blue\n" | |
| "end_header\n" | |
| ) | |
| 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) | |
| 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 | |
| # ===================================================================== | |
| # 3a. Single-image function (original feature) | |
| # ===================================================================== | |
| 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) | |
| # ===================================================================== | |
| 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 | |
| # ===================================================================== | |
| # 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. REST endpoints via Gradio's internal FastAPI app | |
| # ===================================================================== | |
| def status(): | |
| return {"status": "running", "device": str(device)} | |
| 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 = Image.open(io.BytesIO(await file.read())).convert("RGB") | |
| except (UnidentifiedImageError, OSError): | |
| raise HTTPException(status_code=400, detail="Could not decode image.") | |
| try: | |
| 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 | |
| # ===================================================================== | |
| if __name__ == "__main__": | |
| demo.launch() | |