Spaces:
Paused
Paused
| """ | |
| FoundationPose inference server for Hugging Face Spaces with ZeroGPU. | |
| This version uses pure Gradio for ZeroGPU compatibility. | |
| """ | |
| import base64 | |
| import json | |
| import logging | |
| import os | |
| from typing import Dict, List | |
| import cv2 | |
| import gradio as gr | |
| import numpy as np | |
| import spaces | |
| import torch | |
| logging.basicConfig( | |
| level=logging.INFO, | |
| format="[%(asctime)s] %(levelname)s: %(message)s" | |
| ) | |
| logger = logging.getLogger(__name__) | |
| # Check if running in real FoundationPose mode or placeholder mode | |
| USE_REAL_MODEL = os.environ.get("USE_REAL_MODEL", "false").lower() == "true" | |
| class FoundationPoseInference: | |
| """Wrapper for FoundationPose model inference.""" | |
| def __init__(self): | |
| self.model = None | |
| self.device = None | |
| self.initialized = False | |
| self.tracked_objects = {} | |
| self.use_real_model = USE_REAL_MODEL | |
| def initialize_model(self): | |
| """Initialize the FoundationPose model on GPU.""" | |
| if self.initialized: | |
| logger.info("Model already initialized") | |
| return | |
| logger.info("Initializing FoundationPose model...") | |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| logger.info(f"Using device: {self.device}") | |
| if self.use_real_model: | |
| try: | |
| logger.info("Loading real FoundationPose model...") | |
| from estimator import FoundationPoseEstimator | |
| self.model = FoundationPoseEstimator( | |
| device=str(self.device), | |
| weights_dir="weights" | |
| ) | |
| logger.info("✓ Real FoundationPose model initialized successfully") | |
| except Exception as e: | |
| logger.error(f"Failed to initialize real model: {e}", exc_info=True) | |
| logger.warning("Falling back to placeholder mode") | |
| self.use_real_model = False | |
| self.model = None | |
| else: | |
| logger.info("Using placeholder mode (set USE_REAL_MODEL=true for real inference)") | |
| self.model = None | |
| self.initialized = True | |
| logger.info("FoundationPose inference ready") | |
| def register_object( | |
| self, | |
| object_id: str, | |
| reference_images: List[np.ndarray], | |
| camera_intrinsics: Dict = None, | |
| mesh_path: str = None | |
| ) -> bool: | |
| """Register an object for tracking with reference images.""" | |
| if not self.initialized: | |
| self.initialize_model() | |
| logger.info(f"Registering object '{object_id}' with {len(reference_images)} reference images") | |
| if self.use_real_model and self.model is not None: | |
| try: | |
| success = self.model.register_object( | |
| object_id=object_id, | |
| reference_images=reference_images, | |
| camera_intrinsics=camera_intrinsics, | |
| mesh_path=mesh_path | |
| ) | |
| if success: | |
| self.tracked_objects[object_id] = { | |
| "num_references": len(reference_images), | |
| "camera_intrinsics": camera_intrinsics, | |
| "mesh_path": mesh_path | |
| } | |
| return success | |
| except Exception as e: | |
| logger.error(f"Registration failed: {e}", exc_info=True) | |
| return False | |
| else: | |
| self.tracked_objects[object_id] = { | |
| "num_references": len(reference_images), | |
| "camera_intrinsics": camera_intrinsics, | |
| "mesh_path": mesh_path | |
| } | |
| logger.info(f"✓ Object '{object_id}' registered (placeholder mode)") | |
| return True | |
| def estimate_pose( | |
| self, | |
| object_id: str, | |
| query_image: np.ndarray, | |
| camera_intrinsics: Dict = None, | |
| depth_image: np.ndarray = None, | |
| mask: np.ndarray = None | |
| ) -> Dict: | |
| """Estimate 6D pose of an object in a query image.""" | |
| if not self.initialized: | |
| return {"success": False, "error": "Model not initialized"} | |
| if object_id not in self.tracked_objects: | |
| return {"success": False, "error": f"Object '{object_id}' not registered"} | |
| logger.info(f"Estimating pose for object '{object_id}'") | |
| if self.use_real_model and self.model is not None: | |
| try: | |
| pose_result = self.model.estimate_pose( | |
| object_id=object_id, | |
| rgb_image=query_image, | |
| depth_image=depth_image, | |
| mask=mask, | |
| camera_intrinsics=camera_intrinsics | |
| ) | |
| if pose_result is None: | |
| return { | |
| "success": False, | |
| "error": "Pose estimation returned None", | |
| "poses": [] | |
| } | |
| return { | |
| "success": True, | |
| "poses": [pose_result] | |
| } | |
| except Exception as e: | |
| logger.error(f"Pose estimation error: {e}", exc_info=True) | |
| return {"success": False, "error": str(e), "poses": []} | |
| else: | |
| logger.info("Placeholder mode: returning empty pose result") | |
| return { | |
| "success": True, | |
| "poses": [], | |
| "note": "Placeholder mode - set USE_REAL_MODEL=true for real inference" | |
| } | |
| # Global model instance | |
| pose_estimator = FoundationPoseInference() | |
| # Gradio wrapper functions with @spaces.GPU decorators | |
| def gradio_initialize(object_id: str, reference_files: List, fx: float, fy: float, cx: float, cy: float): | |
| """Gradio wrapper for object initialization.""" | |
| try: | |
| if not reference_files: | |
| return "Error: No reference images provided" | |
| # Load reference images | |
| reference_images = [] | |
| for file in reference_files: | |
| img = cv2.imread(file.name) | |
| if img is None: | |
| continue | |
| img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) | |
| reference_images.append(img) | |
| if not reference_images: | |
| return "Error: Could not load any reference images" | |
| # Prepare camera intrinsics | |
| camera_intrinsics = { | |
| "fx": fx, | |
| "fy": fy, | |
| "cx": cx, | |
| "cy": cy | |
| } | |
| # Register object | |
| success = pose_estimator.register_object( | |
| object_id=object_id, | |
| reference_images=reference_images, | |
| camera_intrinsics=camera_intrinsics | |
| ) | |
| if success: | |
| return f"✓ Object '{object_id}' initialized with {len(reference_images)} reference images" | |
| else: | |
| return f"✗ Failed to initialize object '{object_id}'" | |
| except Exception as e: | |
| logger.error(f"Gradio initialization error: {e}", exc_info=True) | |
| return f"Error: {str(e)}" | |
| def gradio_estimate(object_id: str, query_image: np.ndarray, fx: float, fy: float, cx: float, cy: float): | |
| """Gradio wrapper for pose estimation.""" | |
| try: | |
| if query_image is None: | |
| return "Error: No query image provided", None | |
| # Prepare camera intrinsics | |
| camera_intrinsics = { | |
| "fx": fx, | |
| "fy": fy, | |
| "cx": cx, | |
| "cy": cy | |
| } | |
| # Estimate pose | |
| result = pose_estimator.estimate_pose( | |
| object_id=object_id, | |
| query_image=query_image, | |
| camera_intrinsics=camera_intrinsics | |
| ) | |
| if not result.get("success"): | |
| error = result.get("error", "Unknown error") | |
| return f"✗ Estimation failed: {error}", None | |
| poses = result.get("poses", []) | |
| note = result.get("note", "") | |
| # Format output | |
| if not poses: | |
| output = "⚠ No poses detected\n" | |
| if note: | |
| output += f"\nNote: {note}" | |
| return output, query_image | |
| output = f"✓ Detected {len(poses)} pose(s):\n\n" | |
| for i, pose in enumerate(poses): | |
| output += f"Pose {i + 1}:\n" | |
| output += f" Object ID: {pose.get('object_id', 'unknown')}\n" | |
| if 'position' in pose: | |
| pos = pose['position'] | |
| output += f" Position:\n" | |
| output += f" x: {pos.get('x', 0):.4f} m\n" | |
| output += f" y: {pos.get('y', 0):.4f} m\n" | |
| output += f" z: {pos.get('z', 0):.4f} m\n" | |
| if 'orientation' in pose: | |
| ori = pose['orientation'] | |
| output += f" Orientation (quaternion):\n" | |
| output += f" w: {ori.get('w', 0):.4f}\n" | |
| output += f" x: {ori.get('x', 0):.4f}\n" | |
| output += f" y: {ori.get('y', 0):.4f}\n" | |
| output += f" z: {ori.get('z', 0):.4f}\n" | |
| if 'confidence' in pose: | |
| output += f" Confidence: {pose['confidence']:.2%}\n" | |
| output += "\n" | |
| return output, query_image | |
| except Exception as e: | |
| logger.error(f"Gradio estimation error: {e}", exc_info=True) | |
| return f"Error: {str(e)}", None | |
| # Gradio UI | |
| with gr.Blocks(title="FoundationPose Inference", theme=gr.themes.Soft()) as demo: | |
| gr.Markdown("# 🎯 FoundationPose 6D Object Pose Estimation") | |
| mode_indicator = gr.Markdown( | |
| f"**Mode:** {'🟢 Real FoundationPose' if USE_REAL_MODEL else '🟡 Placeholder'}", | |
| elem_id="mode" | |
| ) | |
| with gr.Tabs(): | |
| # Tab 1: Initialize Object | |
| with gr.Tab("Initialize Object"): | |
| gr.Markdown(""" | |
| Upload reference images of your object from different angles (8-20 images recommended). | |
| The model will learn the object's appearance for pose estimation. | |
| """) | |
| with gr.Row(): | |
| with gr.Column(): | |
| init_object_id = gr.Textbox( | |
| label="Object ID", | |
| placeholder="e.g., target_cube", | |
| value="target_cube" | |
| ) | |
| init_ref_files = gr.File( | |
| label="Reference Images", | |
| file_count="multiple", | |
| file_types=["image"] | |
| ) | |
| gr.Markdown("### Camera Intrinsics") | |
| with gr.Row(): | |
| init_fx = gr.Number(label="fx (focal length x)", value=500.0) | |
| init_fy = gr.Number(label="fy (focal length y)", value=500.0) | |
| with gr.Row(): | |
| init_cx = gr.Number(label="cx (principal point x)", value=320.0) | |
| init_cy = gr.Number(label="cy (principal point y)", value=240.0) | |
| init_button = gr.Button("Initialize Object", variant="primary") | |
| with gr.Column(): | |
| init_output = gr.Textbox( | |
| label="Initialization Result", | |
| lines=5, | |
| interactive=False | |
| ) | |
| init_button.click( | |
| fn=gradio_initialize, | |
| inputs=[init_object_id, init_ref_files, init_fx, init_fy, init_cx, init_cy], | |
| outputs=init_output | |
| ) | |
| # Tab 2: Estimate Pose | |
| with gr.Tab("Estimate Pose"): | |
| gr.Markdown(""" | |
| Upload a query image containing the initialized object. | |
| The model will estimate the 6D pose (position + orientation). | |
| """) | |
| with gr.Row(): | |
| with gr.Column(): | |
| est_object_id = gr.Textbox( | |
| label="Object ID", | |
| placeholder="e.g., target_cube", | |
| value="target_cube" | |
| ) | |
| est_query_image = gr.Image( | |
| label="Query Image", | |
| type="numpy" | |
| ) | |
| gr.Markdown("### Camera Intrinsics") | |
| with gr.Row(): | |
| est_fx = gr.Number(label="fx (focal length x)", value=500.0) | |
| est_fy = gr.Number(label="fy (focal length y)", value=500.0) | |
| with gr.Row(): | |
| est_cx = gr.Number(label="cx (principal point x)", value=320.0) | |
| est_cy = gr.Number(label="cy (principal point y)", value=240.0) | |
| est_button = gr.Button("Estimate Pose", variant="primary") | |
| with gr.Column(): | |
| est_output = gr.Textbox( | |
| label="Pose Estimation Result", | |
| lines=15, | |
| interactive=False | |
| ) | |
| est_viz = gr.Image(label="Query Image") | |
| est_button.click( | |
| fn=gradio_estimate, | |
| inputs=[est_object_id, est_query_image, est_fx, est_fy, est_cx, est_cy], | |
| outputs=[est_output, est_viz] | |
| ) | |
| gr.Markdown(""" | |
| --- | |
| ## API Documentation | |
| This Space uses Gradio's built-in API. For programmatic access, use the `gradio_client` library: | |
| ```python | |
| from gradio_client import Client | |
| client = Client("https://gpue-foundationpose.hf.space") | |
| # Initialize object | |
| result = client.predict( | |
| object_id="target_cube", | |
| reference_files=[file1, file2, ...], | |
| fx=500.0, fy=500.0, cx=320.0, cy=240.0, | |
| api_name="/gradio_initialize" | |
| ) | |
| # Estimate pose | |
| result = client.predict( | |
| object_id="target_cube", | |
| query_image=image, | |
| fx=500.0, fy=500.0, cx=320.0, cy=240.0, | |
| api_name="/gradio_estimate" | |
| ) | |
| ``` | |
| See [client.py](https://huggingface.co/spaces/gpue/foundationpose/blob/main/client.py) for a complete example. | |
| """) | |
| if __name__ == "__main__": | |
| demo.launch() | |