app.py

"""
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 pathlib import Path
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"


def download_weights_if_needed():
    """Download model weights from HuggingFace model repository if needed."""
    weights_dir = Path("weights")

    # Check if weights already exist
    if weights_dir.exists() and any(weights_dir.glob("**/*.pth")):
        logger.info("Model weights already exist locally")
        return True

    # Only download if USE_REAL_MODEL is enabled
    if not USE_REAL_MODEL:
        logger.info("Placeholder mode - skipping weight download")
        return False

    try:
        from huggingface_hub import snapshot_download

        model_repo = os.environ.get("FOUNDATIONPOSE_MODEL_REPO", "gpue/foundationpose-weights")
        hf_token = os.environ.get("HF_TOKEN")

        logger.info(f"Downloading model weights from {model_repo}...")

        snapshot_download(
            repo_id=model_repo,
            local_dir=str(weights_dir),
            token=hf_token,
            repo_type="model"
        )

        logger.info("✓ Model weights downloaded successfully")
        return True

    except Exception as e:
        logger.error(f"Failed to download weights: {e}")
        logger.warning("Falling back to placeholder mode")
        return False


# Download weights at startup if needed
logger.info("Checking model weights...")
download_weights_if_needed()


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
@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 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()