foodforthought-cli 0.2.7__py3-none-any.whl → 0.3.0__py3-none-any.whl

ate/urdf/models/__init__.py

@@ -0,0 +1,110 @@
+"""
+Model loading and caching for URDF scan pipeline.
+
+Provides lazy loading and caching for large ML models:
+- SAM 2 (Segment Anything Model 2) for temporal segmentation
+- Depth Anything V2 for metric depth estimation
+"""
+
+from typing import Optional
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class ModelCache:
+    """
+    Singleton cache for expensive ML models.
+
+    Models are loaded lazily on first access and cached for reuse.
+    This prevents redundant loading when processing multiple sessions.
+    """
+
+    _sam2_model = None
+    _sam2_predictor = None
+    _depth_model = None
+    _device: str = "cpu"
+
+    @classmethod
+    def set_device(cls, device: str) -> None:
+        """Set the device for model inference (cuda or cpu)."""
+        cls._device = device
+        # Clear cached models if device changes
+        if cls._sam2_model is not None or cls._depth_model is not None:
+            logger.info(f"Device changed to {device}, clearing model cache")
+            cls.clear()
+
+    @classmethod
+    def get_device(cls) -> str:
+        """Get the current device setting."""
+        return cls._device
+
+    @classmethod
+    def get_sam2(cls):
+        """
+        Get SAM 2 model, loading if necessary.
+
+        Returns:
+            SAM 2 predictor ready for video inference
+        """
+        if cls._sam2_predictor is None:
+            from .sam2 import load_sam2_predictor
+            logger.info(f"Loading SAM 2 model on {cls._device}...")
+            cls._sam2_predictor = load_sam2_predictor(cls._device)
+            logger.info("SAM 2 model loaded")
+        return cls._sam2_predictor
+
+    @classmethod
+    def get_depth_model(cls):
+        """
+        Get Depth Anything V2 model, loading if necessary.
+
+        Returns:
+            Depth Anything V2 model ready for inference
+        """
+        if cls._depth_model is None:
+            from .depth_anything import load_depth_model
+            logger.info(f"Loading Depth Anything V2 on {cls._device}...")
+            cls._depth_model = load_depth_model(cls._device)
+            logger.info("Depth Anything V2 loaded")
+        return cls._depth_model
+
+    @classmethod
+    def clear(cls) -> None:
+        """Clear all cached models to free memory."""
+        cls._sam2_model = None
+        cls._sam2_predictor = None
+        cls._depth_model = None
+        logger.info("Model cache cleared")
+
+    @classmethod
+    def is_loaded(cls, model_name: str) -> bool:
+        """Check if a specific model is loaded."""
+        if model_name == "sam2":
+            return cls._sam2_predictor is not None
+        elif model_name == "depth":
+            return cls._depth_model is not None
+        return False
+
+
+def check_cuda_available() -> bool:
+    """Check if CUDA is available for GPU acceleration."""
+    try:
+        import torch
+        return torch.cuda.is_available()
+    except ImportError:
+        return False
+
+
+def get_recommended_device() -> str:
+    """Get the recommended device based on available hardware."""
+    if check_cuda_available():
+        return "cuda"
+    return "cpu"
+
+
+__all__ = [
+    "ModelCache",
+    "check_cuda_available",
+    "get_recommended_device",
+]

ate/urdf/models/depth_anything.py

@@ -0,0 +1,253 @@
+"""
+Depth Anything V2 wrapper for metric depth estimation.
+
+This module provides:
+- load_depth_model: Load the Depth Anything V2 model
+- DepthEstimator: Wrapper for depth inference
+
+When Depth Anything V2 is not installed, a mock implementation
+is used that generates synthetic depth maps for testing.
+
+Installation:
+    pip install depth-anything-v2
+
+References:
+    https://github.com/DepthAnything/Depth-Anything-V2
+"""
+
+import logging
+from typing import Optional, Tuple
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+try:
+    import numpy as np
+    NUMPY_AVAILABLE = True
+except ImportError:
+    NUMPY_AVAILABLE = False
+    np = None
+
+try:
+    import torch
+    import torch.nn.functional as F
+    from depth_anything_v2.dpt import DepthAnythingV2
+    DEPTH_ANYTHING_AVAILABLE = True
+except ImportError:
+    DEPTH_ANYTHING_AVAILABLE = False
+    torch = None
+
+
+class DepthEstimator:
+    """
+    Wrapper for Depth Anything V2 model.
+
+    Provides metric depth estimation from single RGB images.
+    """
+
+    def __init__(self, model, device: str = "cpu"):
+        """
+        Initialize estimator wrapper.
+
+        Args:
+            model: Depth Anything V2 model instance
+            device: Compute device
+        """
+        self.model = model
+        self.device = device
+
+    def estimate(self, image: "np.ndarray") -> "np.ndarray":
+        """
+        Estimate depth from an RGB image.
+
+        Args:
+            image: BGR image from OpenCV (H, W, 3)
+
+        Returns:
+            Depth map (H, W) in relative units
+        """
+        import cv2
+
+        # Convert BGR to RGB
+        rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+        # Prepare input
+        h, w = rgb.shape[:2]
+
+        # Resize to model input size (typically 518x518 for ViT)
+        input_size = 518
+        rgb_resized = cv2.resize(rgb, (input_size, input_size))
+
+        # Normalize
+        rgb_tensor = torch.from_numpy(rgb_resized).float().permute(2, 0, 1)
+        rgb_tensor = rgb_tensor / 255.0
+        rgb_tensor = rgb_tensor.unsqueeze(0).to(self.device)
+
+        # Inference
+        with torch.no_grad():
+            depth = self.model(rgb_tensor)
+
+        # Resize back to original
+        depth = F.interpolate(
+            depth.unsqueeze(1),
+            size=(h, w),
+            mode="bilinear",
+            align_corners=False,
+        )
+
+        depth_np = depth.squeeze().cpu().numpy()
+        return depth_np
+
+    def estimate_metric(
+        self,
+        image: "np.ndarray",
+        scale_factor: float = 1.0,
+    ) -> "np.ndarray":
+        """
+        Estimate metric depth with scale correction.
+
+        Args:
+            image: BGR image from OpenCV
+            scale_factor: Scale factor to convert to meters
+
+        Returns:
+            Depth map (H, W) in meters
+        """
+        depth = self.estimate(image)
+        return depth * scale_factor
+
+
+class MockDepthEstimator:
+    """
+    Mock depth estimator for testing without the real model.
+
+    Generates synthetic depth maps based on image brightness,
+    simulating closer objects being brighter (like typical indoor lighting).
+    """
+
+    def __init__(self, device: str = "cpu"):
+        """Initialize mock estimator."""
+        self.device = device
+
+    def estimate(self, image: "np.ndarray") -> "np.ndarray":
+        """
+        Generate mock depth from image brightness.
+
+        Args:
+            image: BGR image
+
+        Returns:
+            Mock depth map
+        """
+        import cv2
+
+        # Convert to grayscale
+        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY).astype(np.float32)
+
+        # Invert: brighter = closer = smaller depth
+        depth = 255.0 - gray
+
+        # Normalize to [0.5, 5.0] meters range (typical indoor scene)
+        depth = 0.5 + (depth / 255.0) * 4.5
+
+        # Add some noise for realism
+        noise = np.random.normal(0, 0.1, depth.shape).astype(np.float32)
+        depth = depth + noise
+        depth = np.clip(depth, 0.3, 6.0)
+
+        return depth
+
+    def estimate_metric(
+        self,
+        image: "np.ndarray",
+        scale_factor: float = 1.0,
+    ) -> "np.ndarray":
+        """Generate mock metric depth."""
+        depth = self.estimate(image)
+        # Scale factor adjusts the overall scale
+        return depth * scale_factor
+
+
+def load_depth_model(device: str = "cpu") -> DepthEstimator:
+    """
+    Load Depth Anything V2 model.
+
+    Args:
+        device: Compute device ("cuda" or "cpu")
+
+    Returns:
+        DepthEstimator instance (or mock if unavailable)
+    """
+    if DEPTH_ANYTHING_AVAILABLE:
+        try:
+            logger.info("Loading Depth Anything V2 model...")
+
+            # Model configurations
+            model_configs = {
+                "vits": {"encoder": "vits", "features": 64, "out_channels": [48, 96, 192, 384]},
+                "vitb": {"encoder": "vitb", "features": 128, "out_channels": [96, 192, 384, 768]},
+                "vitl": {"encoder": "vitl", "features": 256, "out_channels": [256, 512, 1024, 1024]},
+            }
+
+            # Try to find checkpoint
+            checkpoint_paths = [
+                Path.home() / ".cache" / "depth_anything_v2" / "depth_anything_v2_vitl.pth",
+                Path.home() / ".cache" / "depth_anything_v2" / "depth_anything_v2_vitb.pth",
+                Path.home() / ".cache" / "depth_anything_v2" / "depth_anything_v2_vits.pth",
+                Path("./checkpoints/depth_anything_v2_vitb.pth"),
+            ]
+
+            checkpoint = None
+            config_name = None
+            for path in checkpoint_paths:
+                if path.exists():
+                    checkpoint = str(path)
+                    # Infer config from filename
+                    if "vitl" in path.name:
+                        config_name = "vitl"
+                    elif "vitb" in path.name:
+                        config_name = "vitb"
+                    else:
+                        config_name = "vits"
+                    break
+
+            if checkpoint is None:
+                logger.warning(
+                    "Depth Anything V2 checkpoint not found. Using mock estimator. "
+                    "Download from: https://github.com/DepthAnything/Depth-Anything-V2"
+                )
+                return MockDepthEstimator(device)
+
+            # Build model
+            config = model_configs[config_name]
+            model = DepthAnythingV2(**config)
+            model.load_state_dict(torch.load(checkpoint, map_location=device))
+            model.to(device)
+            model.eval()
+
+            logger.info(f"Loaded Depth Anything V2 ({config_name}) on {device}")
+            return DepthEstimator(model, device)
+
+        except Exception as e:
+            logger.warning(f"Failed to load Depth Anything V2: {e}. Using mock estimator.")
+            return MockDepthEstimator(device)
+
+    else:
+        logger.warning(
+            "Depth Anything V2 not installed. Using mock estimator. "
+            "Install with: pip install depth-anything-v2"
+        )
+        return MockDepthEstimator(device)
+
+
+def check_depth_anything_available() -> bool:
+    """Check if Depth Anything V2 is available."""
+    return DEPTH_ANYTHING_AVAILABLE
+
+
+__all__ = [
+    "DepthEstimator",
+    "MockDepthEstimator",
+    "load_depth_model",
+    "check_depth_anything_available",
+]

ate/urdf/models/sam2.py

@@ -0,0 +1,324 @@
+"""
+SAM 2 (Segment Anything Model 2) wrapper for temporal segmentation.
+
+This module provides:
+- load_sam2_predictor: Load the SAM 2 video predictor
+- SAM2VideoPredictor: Wrapper for video segmentation
+
+When SAM 2 is not installed, a mock implementation is used that
+generates simple masks based on color clustering for testing.
+
+Installation:
+    pip install segment-anything-2
+
+References:
+    https://github.com/facebookresearch/segment-anything-2
+"""
+
+import logging
+from typing import Dict, List, Optional, Tuple
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+try:
+    import numpy as np
+    NUMPY_AVAILABLE = True
+except ImportError:
+    NUMPY_AVAILABLE = False
+    np = None
+
+try:
+    import torch
+    from sam2.build_sam import build_sam2_video_predictor
+    SAM2_AVAILABLE = True
+except ImportError:
+    SAM2_AVAILABLE = False
+    torch = None
+
+
+class SAM2VideoPredictor:
+    """
+    Wrapper for SAM 2 video segmentation.
+
+    Provides temporal mask propagation from initial point prompts.
+    """
+
+    def __init__(self, predictor, device: str = "cpu"):
+        """
+        Initialize predictor wrapper.
+
+        Args:
+            predictor: SAM 2 video predictor instance
+            device: Compute device
+        """
+        self.predictor = predictor
+        self.device = device
+        self._state = None
+
+    def initialize(self, video_path: str) -> None:
+        """
+        Initialize predictor with a video.
+
+        Args:
+            video_path: Path to video file
+        """
+        import cv2
+        import tempfile
+        import shutil
+        from pathlib import Path
+
+        # Extract frames using OpenCV (works on Apple Silicon)
+        # SAM 2's init_state can use an image directory instead of video
+        self._temp_dir = tempfile.mkdtemp(prefix="sam2_frames_")
+        
+        cap = cv2.VideoCapture(video_path)
+        frame_idx = 0
+        while True:
+            ret, frame = cap.read()
+            if not ret:
+                break
+            # Save as JPEG with zero-padded frame numbers
+            frame_path = Path(self._temp_dir) / f"{frame_idx:06d}.jpg"
+            cv2.imwrite(str(frame_path), frame)
+            frame_idx += 1
+        cap.release()
+        
+        logger.info(f"Extracted {frame_idx} frames to {self._temp_dir}")
+        
+        # Use image directory mode instead of video mode
+        self._state = self.predictor.init_state(video_path=self._temp_dir)
+        logger.info(f"SAM 2 initialized with {frame_idx} frames")
+
+    def add_point_prompt(
+        self,
+        frame_idx: int,
+        obj_id: int,
+        point: Tuple[float, float],
+        label: int = 1,
+    ) -> "np.ndarray":
+        """
+        Add a point prompt to initialize tracking.
+
+        Args:
+            frame_idx: Frame index (usually 0)
+            obj_id: Unique object ID for this link
+            point: (x, y) coordinates in frame
+            label: 1 for foreground, 0 for background
+
+        Returns:
+            Initial mask for the object
+        """
+        points = np.array([[point[0], point[1]]], dtype=np.float32)
+        labels = np.array([label], dtype=np.int32)
+
+        _, _, mask_logits = self.predictor.add_new_points_or_box(
+            inference_state=self._state,
+            frame_idx=frame_idx,
+            obj_id=obj_id,
+            points=points,
+            labels=labels,
+        )
+
+        # Convert logits to binary mask
+        mask = (mask_logits > 0).cpu().numpy().squeeze()
+        return mask
+
+    def propagate(self) -> Dict[int, Dict[int, "np.ndarray"]]:
+        """
+        Propagate masks across all video frames.
+
+        Returns:
+            Dict mapping frame_idx -> {obj_id: mask}
+        """
+        results = {}
+
+        for frame_idx, obj_ids, mask_logits in self.predictor.propagate_in_video(
+            self._state
+        ):
+            frame_masks = {}
+            for i, obj_id in enumerate(obj_ids):
+                mask = (mask_logits[i] > 0).cpu().numpy().squeeze()
+                frame_masks[obj_id] = mask
+            results[frame_idx] = frame_masks
+
+        logger.info(f"SAM 2 propagated masks across {len(results)} frames")
+        return results
+
+    def reset(self) -> None:
+        """Reset predictor state."""
+        if self._state is not None:
+            self.predictor.reset_state(self._state)
+        self._state = None
+
+
+class MockSAM2Predictor:
+    """
+    Mock SAM 2 predictor for testing without the real model.
+
+    Uses simple color-based region growing from seed points.
+    This is NOT accurate but allows pipeline testing.
+    """
+
+    def __init__(self, device: str = "cpu"):
+        """Initialize mock predictor."""
+        self.device = device
+        self._video_path: Optional[str] = None
+        self._prompts: Dict[int, Tuple[float, float]] = {}  # obj_id -> point
+        self._frame_count: int = 0
+        self._resolution: Tuple[int, int] = (720, 1280)
+
+    def initialize(self, video_path: str) -> None:
+        """Initialize with video."""
+        import cv2
+
+        self._video_path = video_path
+        cap = cv2.VideoCapture(video_path)
+        self._frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        self._resolution = (height, width)
+        cap.release()
+
+        logger.info(f"Mock SAM 2 initialized: {self._frame_count} frames")
+
+    def add_point_prompt(
+        self,
+        frame_idx: int,
+        obj_id: int,
+        point: Tuple[float, float],
+        label: int = 1,
+    ) -> "np.ndarray":
+        """Add point prompt and return initial mask."""
+        self._prompts[obj_id] = point
+
+        # Generate simple circular mask around point
+        mask = self._generate_circular_mask(point, radius=50)
+        return mask
+
+    def _generate_circular_mask(
+        self,
+        center: Tuple[float, float],
+        radius: int = 50,
+    ) -> "np.ndarray":
+        """Generate a circular mask around a point."""
+        h, w = self._resolution
+        mask = np.zeros((h, w), dtype=bool)
+
+        cx, cy = int(center[0]), int(center[1])
+        y, x = np.ogrid[:h, :w]
+        dist = np.sqrt((x - cx) ** 2 + (y - cy) ** 2)
+        mask[dist <= radius] = True
+
+        return mask
+
+    def propagate(self) -> Dict[int, Dict[int, "np.ndarray"]]:
+        """
+        Generate mock masks for all frames.
+
+        Creates slightly varying masks to simulate motion tracking.
+        """
+        results = {}
+
+        for frame_idx in range(self._frame_count):
+            frame_masks = {}
+
+            for obj_id, point in self._prompts.items():
+                # Add small random offset to simulate motion
+                offset_x = np.random.uniform(-5, 5)
+                offset_y = np.random.uniform(-5, 5)
+                moved_point = (point[0] + offset_x, point[1] + offset_y)
+
+                # Vary radius slightly
+                radius = 50 + np.random.randint(-10, 10)
+                mask = self._generate_circular_mask(moved_point, radius)
+                frame_masks[obj_id] = mask
+
+            results[frame_idx] = frame_masks
+
+        logger.info(f"Mock SAM 2 generated masks for {len(results)} frames")
+        return results
+
+    def reset(self) -> None:
+        """Reset state."""
+        self._prompts = {}
+
+
+def load_sam2_predictor(device: str = "cpu") -> SAM2VideoPredictor:
+    """
+    Load SAM 2 video predictor.
+
+    Args:
+        device: Compute device ("cuda" or "cpu")
+
+    Returns:
+        SAM2VideoPredictor instance (or mock if unavailable)
+    """
+    if SAM2_AVAILABLE:
+        try:
+            # Try to load the real model
+            logger.info("Loading SAM 2 model...")
+
+            # SAM 2 model checkpoint paths
+            # Users should download from:
+            # https://github.com/facebookresearch/segment-anything-2#model-checkpoints
+            checkpoint_paths = [
+                Path.home() / ".cache" / "sam2" / "sam2_hiera_large.pt",
+                Path.home() / ".cache" / "sam2" / "sam2_hiera_base_plus.pt",
+                Path.home() / ".cache" / "sam2" / "sam2_hiera_base.pt",
+                Path("./checkpoints/sam2_hiera_large.pt"),
+                Path("./checkpoints/sam2_hiera_base_plus.pt"),
+            ]
+
+            checkpoint = None
+            for path in checkpoint_paths:
+                if path.exists():
+                    checkpoint = str(path)
+                    break
+
+            if checkpoint is None:
+                logger.warning(
+                    "SAM 2 checkpoint not found. Using mock predictor. "
+                    "Download from: https://github.com/facebookresearch/segment-anything-2"
+                )
+                return MockSAM2Predictor(device)
+
+            # Select config file based on checkpoint
+            if "base_plus" in checkpoint:
+                config_file = "sam2_hiera_b+.yaml"
+            elif "base" in checkpoint:
+                config_file = "sam2_hiera_b.yaml"
+            else:
+                config_file = "sam2_hiera_l.yaml"
+
+            predictor = build_sam2_video_predictor(
+                config_file=config_file,
+                checkpoint=checkpoint,
+                device=device,
+            )
+
+            return SAM2VideoPredictor(predictor, device)
+
+        except Exception as e:
+            logger.warning(f"Failed to load SAM 2: {e}. Using mock predictor.")
+            return MockSAM2Predictor(device)
+
+    else:
+        logger.warning(
+            "SAM 2 not installed. Using mock predictor. "
+            "Install with: pip install segment-anything-2"
+        )
+        return MockSAM2Predictor(device)
+
+
+def check_sam2_available() -> bool:
+    """Check if SAM 2 is available."""
+    return SAM2_AVAILABLE
+
+
+__all__ = [
+    "SAM2VideoPredictor",
+    "MockSAM2Predictor",
+    "load_sam2_predictor",
+    "check_sam2_available",
+]