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

ate/urdf/depth.py

@@ -0,0 +1,360 @@
+"""
+Depth estimation for URDF generation using Depth Anything V2.
+
+This module handles Phase 2b of the pipeline:
+1. Load Depth Anything V2 model (with caching)
+2. Estimate depth for each video frame
+3. Apply scale factor from reference measurement
+
+Output: Per-frame metric depth maps.
+"""
+
+import logging
+from typing import Dict, Generator, List, Optional, Tuple
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+try:
+    import numpy as np
+    import cv2
+    DEPS_AVAILABLE = True
+except ImportError:
+    DEPS_AVAILABLE = False
+    np = None
+    cv2 = None
+
+
+class DepthError(Exception):
+    """Error during depth estimation."""
+    pass
+
+
+class DepthProcessor:
+    """
+    Depth estimation processor using Depth Anything V2.
+
+    Provides metric depth maps from RGB video frames.
+    """
+
+    def __init__(self, device: str = "cpu"):
+        """
+        Initialize processor.
+
+        Args:
+            device: Compute device ("cuda" or "cpu")
+        """
+        if not DEPS_AVAILABLE:
+            raise DepthError(
+                "Required dependencies not installed. "
+                "Run: pip install numpy opencv-python"
+            )
+
+        self.device = device
+        self._estimator = None
+
+    def _get_estimator(self):
+        """Get or load the depth estimator."""
+        if self._estimator is None:
+            from .models import ModelCache
+            ModelCache.set_device(self.device)
+            self._estimator = ModelCache.get_depth_model()
+        return self._estimator
+
+    def estimate_frame(
+        self,
+        frame: "np.ndarray",
+        scale_factor: float = 1.0,
+    ) -> "np.ndarray":
+        """
+        Estimate metric depth for a single frame.
+
+        Args:
+            frame: BGR image from OpenCV
+            scale_factor: Scale factor to convert to meters
+
+        Returns:
+            Depth map (H, W) in meters
+        """
+        estimator = self._get_estimator()
+        return estimator.estimate_metric(frame, scale_factor)
+
+    def process_video(
+        self,
+        video_path: str,
+        scale_factor: float = 1.0,
+        frame_skip: int = 1,
+        progress_callback: Optional[callable] = None,
+    ) -> Generator[Tuple[int, "np.ndarray", "np.ndarray"], None, None]:
+        """
+        Process video frames and yield depth maps.
+
+        Args:
+            video_path: Path to video file
+            scale_factor: Scale factor for metric depth
+            frame_skip: Process every Nth frame (1 = all frames)
+            progress_callback: Optional callback(current, total)
+
+        Yields:
+            Tuple of (frame_idx, rgb_frame, depth_map)
+        """
+        cap = cv2.VideoCapture(video_path)
+        if not cap.isOpened():
+            raise DepthError(f"Could not open video: {video_path}")
+
+        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        frame_idx = 0
+
+        try:
+            while True:
+                ret, frame = cap.read()
+                if not ret:
+                    break
+
+                if frame_idx % frame_skip == 0:
+                    depth = self.estimate_frame(frame, scale_factor)
+                    yield frame_idx, frame, depth
+
+                    if progress_callback:
+                        progress_callback(frame_idx + 1, total_frames)
+
+                frame_idx += 1
+
+        finally:
+            cap.release()
+
+    def process_video_batch(
+        self,
+        video_path: str,
+        scale_factor: float = 1.0,
+        frame_skip: int = 1,
+        progress_callback: Optional[callable] = None,
+    ) -> Dict[int, "np.ndarray"]:
+        """
+        Process all video frames and return depth maps.
+
+        Args:
+            video_path: Path to video file
+            scale_factor: Scale factor for metric depth
+            frame_skip: Process every Nth frame
+            progress_callback: Optional progress callback
+
+        Returns:
+            Dict mapping frame_idx -> depth_map
+        """
+        result = {}
+        for frame_idx, _, depth in self.process_video(
+            video_path, scale_factor, frame_skip, progress_callback
+        ):
+            result[frame_idx] = depth
+
+        logger.info(f"Processed {len(result)} depth frames")
+        return result
+
+
+def compute_scale_from_mask(
+    depth: "np.ndarray",
+    mask: "np.ndarray",
+    reference_meters: float,
+    dimension: str = "width",
+) -> float:
+    """
+    Compute scale factor from a masked region and known dimension.
+
+    This implements the Kalib method for scale calibration:
+    Correlate a known physical dimension with its measured size
+    in the depth map.
+
+    Args:
+        depth: Depth map (H, W)
+        mask: Binary mask for the reference object
+        reference_meters: Known physical dimension in meters
+        dimension: "width" or "height" of the masked region
+
+    Returns:
+        Scale factor (meters per depth unit)
+    """
+    if not DEPS_AVAILABLE:
+        raise DepthError("NumPy not available")
+
+    # Get bounding box of mask
+    coords = np.where(mask)
+    if len(coords[0]) == 0:
+        raise DepthError("Empty mask - cannot compute scale")
+
+    y_min, y_max = coords[0].min(), coords[0].max()
+    x_min, x_max = coords[1].min(), coords[1].max()
+
+    # Compute dimension in pixels
+    if dimension == "width":
+        pixel_size = x_max - x_min
+    else:
+        pixel_size = y_max - y_min
+
+    if pixel_size == 0:
+        raise DepthError(f"Zero {dimension} in mask - cannot compute scale")
+
+    # Get median depth in the mask
+    masked_depth = depth[mask]
+    median_depth = np.median(masked_depth)
+
+    if median_depth == 0:
+        raise DepthError("Zero depth in masked region")
+
+    # Compute scale: physical_size / (pixel_size * depth)
+    # This gives us a factor to convert depth * pixels to meters
+    scale = reference_meters / (pixel_size * median_depth)
+
+    logger.info(
+        f"Scale computed: {scale:.6f} m/unit "
+        f"(ref={reference_meters}m, pixels={pixel_size}, depth={median_depth:.2f})"
+    )
+
+    return scale
+
+
+def calibrate_depth_scale(
+    depth_processor: DepthProcessor,
+    video_path: str,
+    masks: Dict[str, Dict[int, "np.ndarray"]],
+    reference_link: str,
+    reference_meters: float,
+    dimension: str = "width",
+) -> float:
+    """
+    Calibrate global depth scale from a reference link.
+
+    Args:
+        depth_processor: DepthProcessor instance
+        video_path: Path to video
+        masks: Segmentation masks from segmentation phase
+        reference_link: Name of the link with known dimension
+        reference_meters: Known dimension in meters
+        dimension: "width" or "height"
+
+    Returns:
+        Calibrated scale factor
+    """
+    if reference_link not in masks:
+        raise DepthError(
+            f"Reference link '{reference_link}' not found in masks. "
+            f"Available: {list(masks.keys())}"
+        )
+
+    # Get first frame's data
+    link_masks = masks[reference_link]
+    first_frame_idx = min(link_masks.keys())
+    mask = link_masks[first_frame_idx]
+
+    # Get depth for first frame
+    cap = cv2.VideoCapture(video_path)
+    cap.set(cv2.CAP_PROP_POS_FRAMES, first_frame_idx)
+    ret, frame = cap.read()
+    cap.release()
+
+    if not ret:
+        raise DepthError(f"Could not read frame {first_frame_idx}")
+
+    # Estimate depth without scaling
+    depth = depth_processor.estimate_frame(frame, scale_factor=1.0)
+
+    # Compute scale
+    scale = compute_scale_from_mask(depth, mask, reference_meters, dimension)
+
+    return scale
+
+
+def visualize_depth(
+    depth: "np.ndarray",
+    colormap: int = cv2.COLORMAP_VIRIDIS,
+    min_depth: Optional[float] = None,
+    max_depth: Optional[float] = None,
+) -> "np.ndarray":
+    """
+    Visualize a depth map as a colored image.
+
+    Args:
+        depth: Depth map (H, W)
+        colormap: OpenCV colormap
+        min_depth: Minimum depth for normalization
+        max_depth: Maximum depth for normalization
+
+    Returns:
+        BGR visualization image
+    """
+    if not DEPS_AVAILABLE:
+        raise DepthError("OpenCV not available")
+
+    # Normalize
+    if min_depth is None:
+        min_depth = depth.min()
+    if max_depth is None:
+        max_depth = depth.max()
+
+    normalized = (depth - min_depth) / (max_depth - min_depth + 1e-8)
+    normalized = np.clip(normalized * 255, 0, 255).astype(np.uint8)
+
+    # Apply colormap
+    colored = cv2.applyColorMap(normalized, colormap)
+
+    return colored
+
+
+def run_depth_estimation(
+    session: "ScanSession",
+    masks: Dict[str, Dict[int, "np.ndarray"]],
+    frame_skip: int = 1,
+    progress_callback: Optional[callable] = None,
+) -> Tuple[Dict[int, "np.ndarray"], float]:
+    """
+    Run depth estimation phase on a scan session.
+
+    Args:
+        session: ScanSession with video
+        masks: Segmentation masks from Phase 2a
+        frame_skip: Process every Nth frame
+        progress_callback: Optional progress callback
+
+    Returns:
+        Tuple of (depth_maps, scale_factor)
+    """
+    # Initialize processor
+    processor = DepthProcessor(device=session.metadata.device)
+
+    # Calibrate scale if reference provided
+    scale_factor = 1.0
+    if session.metadata.scale_ref:
+        from .scale import parse_scale_ref
+        ref_link, ref_meters = parse_scale_ref(session.metadata.scale_ref)
+
+        logger.info(f"Calibrating scale from '{ref_link}' = {ref_meters}m")
+        scale_factor = calibrate_depth_scale(
+            processor,
+            str(session.video_path),
+            masks,
+            ref_link,
+            ref_meters,
+        )
+        session.metadata.scale_factor = scale_factor
+        session.save_metadata()
+
+    # Process all frames
+    depth_maps = processor.process_video_batch(
+        str(session.video_path),
+        scale_factor=scale_factor,
+        frame_skip=frame_skip,
+        progress_callback=progress_callback,
+    )
+
+    logger.info(f"Depth estimation complete: {len(depth_maps)} frames, scale={scale_factor:.6f}")
+    return depth_maps, scale_factor
+
+
+__all__ = [
+    "DepthError",
+    "DepthProcessor",
+    "compute_scale_from_mask",
+    "calibrate_depth_scale",
+    "visualize_depth",
+    "run_depth_estimation",
+]

ate/urdf/inertial.py

@@ -0,0 +1,312 @@
+"""
+Inertial property estimation for URDF generation.
+
+This module handles Phase 5a of the pipeline:
+1. Compute mesh volume from visual/collision meshes
+2. Estimate mass from assumed material density
+3. Calculate center of mass and inertia tensor
+
+Video cannot measure mass directly, so we use geometric estimation
+with configurable material density assumptions.
+"""
+
+import logging
+from typing import Dict, List, Optional, Tuple
+from pathlib import Path
+from dataclasses import dataclass
+
+logger = logging.getLogger(__name__)
+
+try:
+    import numpy as np
+    NUMPY_AVAILABLE = True
+except ImportError:
+    NUMPY_AVAILABLE = False
+    np = None
+
+try:
+    import trimesh
+    TRIMESH_AVAILABLE = True
+except ImportError:
+    TRIMESH_AVAILABLE = False
+    trimesh = None
+
+
+class InertialError(Exception):
+    """Error during inertial estimation."""
+    pass
+
+
+# Common material densities in kg/m^3
+MATERIAL_DENSITIES = {
+    "pla": 1250.0,          # PLA plastic
+    "abs": 1040.0,          # ABS plastic
+    "petg": 1270.0,         # PETG plastic
+    "nylon": 1150.0,        # Nylon
+    "aluminum": 2700.0,     # Aluminum
+    "steel": 7850.0,        # Steel
+    "carbon_fiber": 1600.0, # Carbon fiber composite
+    "wood": 700.0,          # Average wood
+    "foam": 100.0,          # Foam/hollow structures
+    "default": 1200.0,      # Default (generic plastic)
+}
+
+
+@dataclass
+class InertialProperties:
+    """Inertial properties for a single link."""
+    link_name: str
+    mass: float  # kg
+    center_of_mass: List[float]  # [x, y, z] meters
+    inertia: Dict[str, float]  # ixx, ixy, ixz, iyy, iyz, izz (kg*m^2)
+
+    def to_dict(self) -> Dict:
+        return {
+            "link_name": self.link_name,
+            "mass": self.mass,
+            "center_of_mass": self.center_of_mass,
+            "inertia": self.inertia,
+        }
+
+    @property
+    def inertia_tensor(self) -> "np.ndarray":
+        """Get 3x3 inertia tensor matrix."""
+        return np.array([
+            [self.inertia["ixx"], self.inertia["ixy"], self.inertia["ixz"]],
+            [self.inertia["ixy"], self.inertia["iyy"], self.inertia["iyz"]],
+            [self.inertia["ixz"], self.inertia["iyz"], self.inertia["izz"]],
+        ])
+
+
+def compute_mesh_volume(mesh: "trimesh.Trimesh") -> float:
+    """
+    Compute volume of a mesh.
+
+    Args:
+        mesh: Trimesh object (should be watertight)
+
+    Returns:
+        Volume in cubic meters
+    """
+    if not TRIMESH_AVAILABLE:
+        raise InertialError("trimesh not available")
+
+    if not mesh.is_watertight:
+        logger.warning("Mesh is not watertight, volume may be inaccurate")
+        # Attempt to fix
+        mesh.fill_holes()
+
+    return abs(mesh.volume)
+
+
+def compute_inertia_tensor(
+    mesh: "trimesh.Trimesh",
+    density: float,
+) -> Tuple[float, "np.ndarray", "np.ndarray"]:
+    """
+    Compute mass, center of mass, and inertia tensor.
+
+    Uses trimesh's built-in moment_inertia calculation.
+
+    Args:
+        mesh: Trimesh object
+        density: Material density (kg/m^3)
+
+    Returns:
+        Tuple of (mass, center_of_mass, inertia_tensor)
+    """
+    if not TRIMESH_AVAILABLE:
+        raise InertialError("trimesh not available")
+
+    # Compute volume and mass
+    volume = compute_mesh_volume(mesh)
+    mass = volume * density
+
+    # Ensure minimum mass for stability
+    MIN_MASS = 0.001  # 1 gram
+    if mass < MIN_MASS:
+        logger.warning(f"Very small mass {mass:.6f}kg, clamping to {MIN_MASS}kg")
+        mass = MIN_MASS
+
+    # Get center of mass
+    com = mesh.center_mass
+
+    # Get inertia tensor at center of mass
+    # trimesh returns inertia in mesh units, we need to scale by density
+    inertia = mesh.moment_inertia * density
+
+    return mass, com, inertia
+
+
+def estimate_inertial_from_mesh(
+    mesh_path: Path,
+    density: float = 1200.0,
+) -> InertialProperties:
+    """
+    Estimate inertial properties from a mesh file.
+
+    Args:
+        mesh_path: Path to mesh file (OBJ, STL, etc.)
+        density: Material density (kg/m^3)
+
+    Returns:
+        InertialProperties for the mesh
+    """
+    if not TRIMESH_AVAILABLE:
+        raise InertialError("trimesh not available")
+
+    # Load mesh
+    mesh = trimesh.load(str(mesh_path))
+
+    if not isinstance(mesh, trimesh.Trimesh):
+        if hasattr(mesh, 'geometry'):
+            meshes = list(mesh.geometry.values())
+            if meshes:
+                mesh = meshes[0]
+            else:
+                raise InertialError(f"No geometry in {mesh_path}")
+        else:
+            raise InertialError(f"Invalid mesh in {mesh_path}")
+
+    # Compute inertial properties
+    mass, com, inertia = compute_inertia_tensor(mesh, density)
+
+    # Extract link name from path
+    link_name = mesh_path.stem.replace("_visual", "").replace("_collision", "")
+
+    return InertialProperties(
+        link_name=link_name,
+        mass=float(mass),
+        center_of_mass=com.tolist(),
+        inertia={
+            "ixx": float(inertia[0, 0]),
+            "ixy": float(inertia[0, 1]),
+            "ixz": float(inertia[0, 2]),
+            "iyy": float(inertia[1, 1]),
+            "iyz": float(inertia[1, 2]),
+            "izz": float(inertia[2, 2]),
+        },
+    )
+
+
+def estimate_inertial_box(
+    dimensions: List[float],
+    density: float = 1200.0,
+    link_name: str = "link",
+) -> InertialProperties:
+    """
+    Estimate inertial properties for a box shape.
+
+    Useful for simple approximations or fallback.
+
+    Args:
+        dimensions: [width, height, depth] in meters
+        density: Material density
+        link_name: Name for the link
+
+    Returns:
+        InertialProperties
+    """
+    if not NUMPY_AVAILABLE:
+        raise InertialError("NumPy not available")
+
+    w, h, d = dimensions
+    volume = w * h * d
+    mass = volume * density
+
+    # Box inertia formulas
+    ixx = (1/12) * mass * (h**2 + d**2)
+    iyy = (1/12) * mass * (w**2 + d**2)
+    izz = (1/12) * mass * (w**2 + h**2)
+
+    return InertialProperties(
+        link_name=link_name,
+        mass=mass,
+        center_of_mass=[0.0, 0.0, 0.0],
+        inertia={
+            "ixx": ixx, "ixy": 0.0, "ixz": 0.0,
+            "iyy": iyy, "iyz": 0.0,
+            "izz": izz,
+        },
+    )
+
+
+def estimate_all_inertials(
+    session: "ScanSession",
+    density: float = 1200.0,
+    use_collision_mesh: bool = True,
+    progress_callback: Optional[callable] = None,
+) -> Dict[str, InertialProperties]:
+    """
+    Estimate inertial properties for all links.
+
+    Args:
+        session: ScanSession with meshes
+        density: Material density
+        use_collision_mesh: Use collision (simpler) or visual mesh
+        progress_callback: Optional progress callback
+
+    Returns:
+        Dict mapping link_name -> InertialProperties
+    """
+    meshes_dir = session.meshes_dir
+
+    # Find meshes
+    suffix = "_collision.obj" if use_collision_mesh else "_visual.obj"
+    mesh_files = list(meshes_dir.glob(f"*{suffix}"))
+
+    if not mesh_files:
+        # Fall back to other type
+        alt_suffix = "_visual.obj" if use_collision_mesh else "_collision.obj"
+        mesh_files = list(meshes_dir.glob(f"*{alt_suffix}"))
+
+    if not mesh_files:
+        raise InertialError(f"No mesh files found in {meshes_dir}")
+
+    result = {}
+    total = len(mesh_files)
+
+    for i, mesh_path in enumerate(mesh_files):
+        try:
+            props = estimate_inertial_from_mesh(mesh_path, density)
+            result[props.link_name] = props
+
+            logger.info(
+                f"Estimated inertial for '{props.link_name}': "
+                f"mass={props.mass:.4f}kg, com={props.center_of_mass}"
+            )
+        except Exception as e:
+            logger.error(f"Failed to estimate inertial for {mesh_path}: {e}")
+
+        if progress_callback:
+            progress_callback(i + 1, total)
+
+    logger.info(f"Estimated inertials for {len(result)} links")
+    return result
+
+
+def get_density_for_material(material: str) -> float:
+    """
+    Get density for a material name.
+
+    Args:
+        material: Material name (e.g., "pla", "aluminum")
+
+    Returns:
+        Density in kg/m^3
+    """
+    material = material.lower().replace(" ", "_").replace("-", "_")
+    return MATERIAL_DENSITIES.get(material, MATERIAL_DENSITIES["default"])
+
+
+__all__ = [
+    "InertialError",
+    "MATERIAL_DENSITIES",
+    "InertialProperties",
+    "compute_mesh_volume",
+    "compute_inertia_tensor",
+    "estimate_inertial_from_mesh",
+    "estimate_inertial_box",
+    "estimate_all_inertials",
+    "get_density_for_material",
+]