""" Simple FoundationPose API server using FastAPI + Gradio This version uses FastAPI for clean REST API endpoints alongside Gradio UI. """ 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 from fastapi import FastAPI, HTTPException from pydantic import BaseModel 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() # Pydantic models for API class InitializeRequest(BaseModel): object_id: str reference_images_b64: List[str] camera_intrinsics: str = None mesh_path: str = None class EstimateRequest(BaseModel): object_id: str query_image_b64: str camera_intrinsics: str = None depth_image_b64: str = None mask_b64: str = None # Create FastAPI app app = FastAPI() @app.post("/api/initialize") async def api_initialize(request: InitializeRequest): """Initialize object tracking with reference images.""" try: # Decode reference images reference_images = [] for img_b64 in request.reference_images_b64: img_bytes = base64.b64decode(img_b64) img_array = np.frombuffer(img_bytes, dtype=np.uint8) img = cv2.imdecode(img_array, cv2.IMREAD_COLOR) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) reference_images.append(img) # Parse camera intrinsics intrinsics = json.loads(request.camera_intrinsics) if request.camera_intrinsics else None # Register object success = pose_estimator.register_object( object_id=request.object_id, reference_images=reference_images, camera_intrinsics=intrinsics, mesh_path=request.mesh_path ) return { "success": success, "message": f"Object '{request.object_id}' registered with {len(reference_images)} reference images" } except Exception as e: logger.error(f"Initialization error: {e}", exc_info=True) raise HTTPException(status_code=500, detail=str(e)) @app.post("/api/estimate") async def api_estimate(request: EstimateRequest): """Estimate 6D pose from query image.""" try: # Decode query image img_bytes = base64.b64decode(request.query_image_b64) img_array = np.frombuffer(img_bytes, dtype=np.uint8) img = cv2.imdecode(img_array, cv2.IMREAD_COLOR) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # Decode optional depth image depth = None if request.depth_image_b64: depth_bytes = base64.b64decode(request.depth_image_b64) depth = np.frombuffer(depth_bytes, dtype=np.float32) # Decode optional mask mask = None if request.mask_b64: mask_bytes = base64.b64decode(request.mask_b64) mask_array = np.frombuffer(mask_bytes, dtype=np.uint8) mask = cv2.imdecode(mask_array, cv2.IMREAD_GRAYSCALE) # Parse camera intrinsics intrinsics = json.loads(request.camera_intrinsics) if request.camera_intrinsics else None # Estimate pose result = pose_estimator.estimate_pose( object_id=request.object_id, query_image=img, camera_intrinsics=intrinsics, depth_image=depth, mask=mask ) return result except Exception as e: logger.error(f"Estimation error: {e}", exc_info=True) raise HTTPException(status_code=500, detail=str(e)) # Gradio wrapper functions @spaces.GPU(duration=120) 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)}" @spaces.GPU(duration=30) 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 proper @spaces.GPU function calls with gr.Blocks(title="FoundationPose Inference", theme=gr.themes.Soft()) as gradio_app: 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(""" --- ## REST API Endpoints This Space also provides REST API endpoints for programmatic access: - POST `/api/initialize` - Register object with reference images - POST `/api/estimate` - Estimate 6D pose from query image See the [API documentation](https://huggingface.co/spaces/gpue/foundationpose) for details. """) # Mount Gradio to FastAPI app = gr.mount_gradio_app(app, gradio_app, path="/") if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=7860)