estimator.py

"""
FoundationPose Estimator Wrapper

This module wraps the FoundationPose API for easy integration with the Gradio app.
"""

import logging
import sys
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import cv2
import numpy as np
import torch
import trimesh

logger = logging.getLogger(__name__)


class FoundationPoseEstimator:
    """Wrapper for FoundationPose 6D pose estimation."""

    def __init__(self, device: str = "cuda", weights_dir: str = "weights"):
        """Initialize FoundationPose.

        Args:
            device: Device to run inference on ("cuda" or "cpu")
            weights_dir: Path to model weights directory
        """
        self.device = device
        self.weights_dir = Path(weights_dir)

        # Add FoundationPose to Python path
        foundationpose_dir = Path("FoundationPose")
        if foundationpose_dir.exists():
            sys.path.insert(0, str(foundationpose_dir))
        else:
            raise RuntimeError(
                "FoundationPose repository not found. "
                "Clone it with: git clone https://github.com/NVlabs/FoundationPose.git"
            )

        # Import FoundationPose modules
        try:
            from estimater import FoundationPose
            from datareader import SceneReader
            import pytorch3d.transforms as transforms

            self.FoundationPose = FoundationPose
            self.SceneReader = SceneReader
            self.transforms = transforms

        except ImportError as e:
            raise RuntimeError(
                f"Failed to import FoundationPose modules: {e}\n"
                "Make sure FoundationPose is properly installed with all dependencies."
            )

        # Initialize models
        self._init_models()

        # Tracking state
        self.tracked_objects = {}
        self.pose_estimators = {}

    def _init_models(self):
        """Initialize scorer and refiner models."""
        logger.info("Initializing FoundationPose models...")

        try:
            # Load scorer model
            scorer_weights = self.weights_dir / "2024-01-11-20-02-45"
            if not scorer_weights.exists():
                raise FileNotFoundError(f"Scorer weights not found at {scorer_weights}")

            # Load refiner model
            refiner_weights = self.weights_dir / "2023-10-28-18-33-37"
            if not refiner_weights.exists():
                raise FileNotFoundError(f"Refiner weights not found at {refiner_weights}")

            # Import and initialize models (actual implementation depends on FoundationPose API)
            from model import FoundationPoseModel

            self.scorer = FoundationPoseModel(
                checkpoint_dir=str(scorer_weights),
                model_type="scorer"
            ).to(self.device)
            self.scorer.eval()

            self.refiner = FoundationPoseModel(
                checkpoint_dir=str(refiner_weights),
                model_type="refiner"
            ).to(self.device)
            self.refiner.eval()

            # Initialize CUDA rasterization context
            import nvdiffrast.torch as dr
            self.glctx = dr.RasterizeCudaContext()

            logger.info("✓ Models initialized successfully")

        except Exception as e:
            logger.error(f"Failed to initialize models: {e}")
            raise

    def register_object(
        self,
        object_id: str,
        reference_images: List[np.ndarray],
        camera_intrinsics: Optional[Dict] = None,
        mesh_path: Optional[str] = None
    ) -> bool:
        """Register an object for tracking.

        Args:
            object_id: Unique identifier for the object
            reference_images: List of RGB images from different viewpoints
            camera_intrinsics: Camera parameters (fx, fy, cx, cy)
            mesh_path: Optional path to CAD mesh (for model-based mode)

        Returns:
            True if registration successful
        """
        logger.info(f"Registering object '{object_id}'...")

        try:
            # Load or reconstruct mesh
            if mesh_path and Path(mesh_path).exists():
                # Model-based: use CAD mesh
                mesh = trimesh.load(mesh_path)
                logger.info(f"Loaded mesh from {mesh_path}")
            else:
                # Model-free: reconstruct from reference images
                logger.info("Reconstructing mesh from reference images...")
                mesh = self._reconstruct_mesh_from_references(
                    reference_images, camera_intrinsics
                )

            # Create FoundationPose estimator for this object
            estimator = self.FoundationPose(
                model_pts=mesh.vertices,
                model_normals=mesh.vertex_normals,
                mesh=mesh,
                scorer=self.scorer,
                refiner=self.refiner,
                debug_dir=None,
                debug=0,
                glctx=self.glctx
            )

            # Store object data
            self.tracked_objects[object_id] = {
                "mesh": mesh,
                "camera_intrinsics": camera_intrinsics,
                "registered": True
            }
            self.pose_estimators[object_id] = {
                "estimator": estimator,
                "tracking": False,
                "last_pose": None
            }

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

        except Exception as e:
            logger.error(f"Failed to register object: {e}", exc_info=True)
            return False

    def _reconstruct_mesh_from_references(
        self,
        reference_images: List[np.ndarray],
        camera_intrinsics: Optional[Dict]
    ) -> trimesh.Trimesh:
        """Reconstruct 3D mesh from reference images using BundleSDF.

        Args:
            reference_images: List of RGB images
            camera_intrinsics: Camera parameters

        Returns:
            Reconstructed mesh
        """
        # TODO: Implement BundleSDF reconstruction
        # For now, return a simple placeholder mesh
        logger.warning("Mesh reconstruction not fully implemented, using placeholder")

        # Create a simple cube mesh as placeholder
        mesh = trimesh.creation.box(extents=[0.1, 0.1, 0.1])
        return mesh

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

        Args:
            object_id: ID of registered object
            rgb_image: RGB image (H, W, 3)
            depth_image: Optional depth map (H, W)
            mask: Optional object segmentation mask (H, W)
            camera_intrinsics: Camera parameters

        Returns:
            Pose dictionary with position, orientation, and confidence
        """
        if object_id not in self.pose_estimators:
            logger.error(f"Object '{object_id}' not registered")
            return None

        try:
            estimator_data = self.pose_estimators[object_id]
            estimator = estimator_data["estimator"]

            # Get camera intrinsics
            if camera_intrinsics is None:
                camera_intrinsics = self.tracked_objects[object_id]["camera_intrinsics"]

            K = self._build_intrinsics_matrix(camera_intrinsics, rgb_image.shape)

            # Generate synthetic depth if not provided
            if depth_image is None:
                depth_image = np.zeros((rgb_image.shape[0], rgb_image.shape[1]), dtype=np.float32)

            # Auto-segment if mask not provided
            if mask is None:
                mask = self._segment_object(rgb_image)

            # First frame: register
            if not estimator_data["tracking"]:
                logger.info(f"Initial registration for '{object_id}'")
                pose = estimator.register(
                    K=K,
                    rgb=rgb_image,
                    depth=depth_image,
                    ob_mask=mask,
                    iteration=5  # Number of refinement iterations
                )
                estimator_data["tracking"] = True
                estimator_data["last_pose"] = pose
            else:
                # Subsequent frames: track
                pose = estimator.track_one(
                    rgb=rgb_image,
                    depth=depth_image,
                    K=K,
                    iteration=2
                )
                estimator_data["last_pose"] = pose

            # Convert pose matrix to position + quaternion
            result = self._pose_matrix_to_dict(pose, object_id)

            logger.info(f"Estimated pose for '{object_id}': confidence={result['confidence']:.3f}")
            return result

        except Exception as e:
            logger.error(f"Pose estimation failed: {e}", exc_info=True)
            return None

    def _build_intrinsics_matrix(
        self,
        intrinsics: Optional[Dict],
        image_shape: Tuple[int, int, int]
    ) -> np.ndarray:
        """Build camera intrinsics matrix.

        Args:
            intrinsics: Dict with fx, fy, cx, cy
            image_shape: (H, W, C)

        Returns:
            3x3 intrinsics matrix
        """
        H, W = image_shape[:2]

        if intrinsics:
            fx = intrinsics.get("fx", 500.0)
            fy = intrinsics.get("fy", 500.0)
            cx = intrinsics.get("cx", W / 2)
            cy = intrinsics.get("cy", H / 2)
        else:
            # Default intrinsics
            fx = fy = 500.0
            cx = W / 2
            cy = H / 2

        K = np.array([
            [fx, 0, cx],
            [0, fy, cy],
            [0, 0, 1]
        ], dtype=np.float32)

        return K

    def _segment_object(self, rgb_image: np.ndarray) -> np.ndarray:
        """Segment object from background.

        This is a placeholder - in production, use SAM or similar.

        Args:
            rgb_image: RGB image

        Returns:
            Binary mask
        """
        # Simple color-based segmentation placeholder
        # In production, use Segment Anything Model (SAM)
        H, W = rgb_image.shape[:2]
        mask = np.ones((H, W), dtype=np.uint8) * 255

        logger.warning("Using placeholder segmentation - implement SAM for production")
        return mask

    def _pose_matrix_to_dict(self, pose_matrix: np.ndarray, object_id: str) -> Dict:
        """Convert 4x4 pose matrix to dictionary format.

        Args:
            pose_matrix: 4x4 transformation matrix
            object_id: Object identifier

        Returns:
            Dictionary with position, orientation (quaternion), confidence
        """
        # Extract translation
        position = {
            "x": float(pose_matrix[0, 3]),
            "y": float(pose_matrix[1, 3]),
            "z": float(pose_matrix[2, 3])
        }

        # Extract rotation matrix and convert to quaternion
        rotation_matrix = pose_matrix[:3, :3]
        quat = self._rotation_matrix_to_quaternion(rotation_matrix)

        orientation = {
            "w": float(quat[0]),
            "x": float(quat[1]),
            "y": float(quat[2]),
            "z": float(quat[3])
        }

        # Estimate confidence based on tracking state
        # In production, use actual confidence from the model
        confidence = 0.9 if self.pose_estimators[object_id]["tracking"] else 0.7

        # Get object dimensions from mesh
        mesh = self.tracked_objects[object_id]["mesh"]
        extents = mesh.bounds[1] - mesh.bounds[0]
        dimensions = [float(extents[0]), float(extents[1]), float(extents[2])]

        return {
            "object_id": object_id,
            "position": position,
            "orientation": orientation,
            "confidence": confidence,
            "dimensions": dimensions,
            "timestamp": 0.0  # Add timestamp if needed
        }

    def _rotation_matrix_to_quaternion(self, R: np.ndarray) -> np.ndarray:
        """Convert 3x3 rotation matrix to quaternion (w, x, y, z).

        Args:
            R: 3x3 rotation matrix

        Returns:
            Quaternion as numpy array [w, x, y, z]
        """
        trace = np.trace(R)

        if trace > 0:
            s = 0.5 / np.sqrt(trace + 1.0)
            w = 0.25 / s
            x = (R[2, 1] - R[1, 2]) * s
            y = (R[0, 2] - R[2, 0]) * s
            z = (R[1, 0] - R[0, 1]) * s
        elif R[0, 0] > R[1, 1] and R[0, 0] > R[2, 2]:
            s = 2.0 * np.sqrt(1.0 + R[0, 0] - R[1, 1] - R[2, 2])
            w = (R[2, 1] - R[1, 2]) / s
            x = 0.25 * s
            y = (R[0, 1] + R[1, 0]) / s
            z = (R[0, 2] + R[2, 0]) / s
        elif R[1, 1] > R[2, 2]:
            s = 2.0 * np.sqrt(1.0 + R[1, 1] - R[0, 0] - R[2, 2])
            w = (R[0, 2] - R[2, 0]) / s
            x = (R[0, 1] + R[1, 0]) / s
            y = 0.25 * s
            z = (R[1, 2] + R[2, 1]) / s
        else:
            s = 2.0 * np.sqrt(1.0 + R[2, 2] - R[0, 0] - R[1, 1])
            w = (R[1, 0] - R[0, 1]) / s
            x = (R[0, 2] + R[2, 0]) / s
            y = (R[1, 2] + R[2, 1]) / s
            z = 0.25 * s

        return np.array([w, x, y, z])

    def reset_tracking(self, object_id: str):
        """Reset tracking state for an object.

        Args:
            object_id: Object to reset
        """
        if object_id in self.pose_estimators:
            self.pose_estimators[object_id]["tracking"] = False
            self.pose_estimators[object_id]["last_pose"] = None
            logger.info(f"Reset tracking for '{object_id}'")