app_simple.py

"""
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

    @spaces.GPU(duration=120)  # Allocate GPU for 120 seconds (includes model loading)
    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

    @spaces.GPU(duration=10)
    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 UI (simplified)
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"
    )

    gr.Markdown("""
    API Endpoints:
    - POST `/api/initialize` - Register object
    - POST `/api/estimate` - Estimate pose

    See documentation for usage examples.
    """)


# 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)