client.py

"""
Client for FoundationPose Hugging Face Space API

This client can be used from the robot-ml training pipeline to call the
FoundationPose inference API hosted on Hugging Face Spaces.
"""

import base64
import json
import logging
from io import BytesIO
from pathlib import Path
from typing import Dict, List, Optional

import cv2
import numpy as np
import requests

logger = logging.getLogger(__name__)


class FoundationPoseClient:
    """Client for FoundationPose API."""

    def __init__(self, api_url: str = "https://gpue-foundationpose.hf.space"):
        """Initialize client.

        Args:
            api_url: Base URL of the FoundationPose Space
        """
        self.api_url = api_url.rstrip("/")
        self.session = requests.Session()
        self.session.headers.update({"Content-Type": "application/json"})

    def _encode_image(self, image: np.ndarray) -> str:
        """Encode image as base64 JPEG.

        Args:
            image: RGB image as numpy array

        Returns:
            Base64-encoded JPEG string
        """
        # Convert RGB to BGR for OpenCV
        image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

        # Encode as JPEG
        _, buffer = cv2.imencode(".jpg", image_bgr, [cv2.IMWRITE_JPEG_QUALITY, 85])

        # Convert to base64
        image_b64 = base64.b64encode(buffer).decode("utf-8")
        return image_b64

    def initialize(
        self,
        object_id: str,
        reference_images: List[np.ndarray],
        camera_intrinsics: Optional[Dict] = None
    ) -> bool:
        """Initialize object tracking with reference images.

        Args:
            object_id: Unique ID for the object
            reference_images: List of RGB images (numpy arrays)
            camera_intrinsics: Optional camera parameters

        Returns:
            True if successful

        Raises:
            RuntimeError: If initialization fails
        """
        logger.info(f"Initializing object '{object_id}' with {len(reference_images)} reference images")

        # Encode images
        images_b64 = [self._encode_image(img) for img in reference_images]

        # Prepare request
        payload = {
            "object_id": object_id,
            "reference_images_b64": images_b64,
        }

        if camera_intrinsics:
            payload["camera_intrinsics"] = json.dumps(camera_intrinsics)

        # Send request
        try:
            response = self.session.post(
                f"{self.api_url}/api/initialize",
                json=payload,
                timeout=120  # Long timeout for model loading
            )
            response.raise_for_status()

            result = response.json()

            if not result.get("success"):
                error = result.get("error", "Unknown error")
                raise RuntimeError(f"Initialization failed: {error}")

            logger.info(f"Object '{object_id}' initialized successfully")
            return True

        except requests.exceptions.RequestException as e:
            logger.error(f"API request failed: {e}")
            raise RuntimeError(f"Failed to initialize object: {e}")

    def estimate_pose(
        self,
        object_id: str,
        query_image: np.ndarray,
        camera_intrinsics: Optional[Dict] = None
    ) -> List[Dict]:
        """Estimate 6D pose of object in query image.

        Args:
            object_id: ID of object to detect
            query_image: RGB query image as numpy array
            camera_intrinsics: Optional camera parameters

        Returns:
            List of detected poses:
            [
                {
                    "object_id": str,
                    "position": {"x": float, "y": float, "z": float},
                    "orientation": {"w": float, "x": float, "y": float, "z": float},
                    "confidence": float,
                    "dimensions": [float, float, float]
                }
            ]

        Raises:
            RuntimeError: If estimation fails
        """
        # Encode image
        image_b64 = self._encode_image(query_image)

        # Prepare request
        payload = {
            "object_id": object_id,
            "query_image_b64": image_b64,
        }

        if camera_intrinsics:
            payload["camera_intrinsics"] = json.dumps(camera_intrinsics)

        # Send request
        try:
            response = self.session.post(
                f"{self.api_url}/api/estimate",
                json=payload,
                timeout=30
            )
            response.raise_for_status()

            result = response.json()

            if not result.get("success"):
                error = result.get("error", "Unknown error")
                raise RuntimeError(f"Pose estimation failed: {error}")

            return result.get("poses", [])

        except requests.exceptions.RequestException as e:
            logger.error(f"API request failed: {e}")
            raise RuntimeError(f"Failed to estimate pose: {e}")


def load_reference_images(directory: Path) -> List[np.ndarray]:
    """Load reference images from directory.

    Args:
        directory: Path to directory containing images

    Returns:
        List of RGB images as numpy arrays
    """
    images = []
    for img_path in sorted(directory.glob("*.jpg")):
        img = cv2.imread(str(img_path))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        images.append(img)

    logger.info(f"Loaded {len(images)} reference images from {directory}")
    return images


# Example usage
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)

    # Initialize client
    client = FoundationPoseClient()

    # Load reference images
    ref_dir = Path("../training/perception/reference/target_cube")
    if ref_dir.exists():
        ref_images = load_reference_images(ref_dir)

        # Initialize object
        client.initialize("target_cube", ref_images)

        # Estimate pose on first reference image (for testing)
        poses = client.estimate_pose("target_cube", ref_images[0])
        print(f"Detected {len(poses)} poses:")
        for pose in poses:
            print(f"  {pose}")
    else:
        print(f"Reference directory not found: {ref_dir}")
        print("Run 'make capture-reference' to collect reference images first")