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

ate/urdf/collision.py

@@ -0,0 +1,271 @@
+"""
+Collision mesh generation via convex decomposition.
+
+This module handles Phase 4b of the pipeline:
+1. Load visual meshes
+2. Decompose into convex hulls using CoACD or V-HACD
+3. Save collision meshes
+
+Physics engines require convex collision shapes for stable simulation.
+"""
+
+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 trimesh
+    TRIMESH_AVAILABLE = True
+except ImportError:
+    TRIMESH_AVAILABLE = False
+    trimesh = None
+
+
+class CollisionError(Exception):
+    """Error during collision mesh generation."""
+    pass
+
+
+def coacd_decompose(
+    mesh: "trimesh.Trimesh",
+    max_hulls: int = 8,
+    threshold: float = 0.05,
+) -> List["trimesh.Trimesh"]:
+    """
+    Decompose mesh into convex hulls using CoACD.
+
+    CoACD (Approximate Convex Decomposition) is more accurate than V-HACD
+    but may not be available on all systems.
+
+    Args:
+        mesh: Input mesh to decompose
+        max_hulls: Maximum number of convex hulls
+        threshold: Concavity threshold
+
+    Returns:
+        List of convex hull meshes
+    """
+    try:
+        import coacd
+
+        # CoACD expects vertices and faces
+        vertices = np.array(mesh.vertices, dtype=np.float64)
+        faces = np.array(mesh.faces, dtype=np.int32)
+
+        # Run decomposition
+        parts = coacd.run_coacd(
+            vertices,
+            faces,
+            threshold=threshold,
+            max_convex_hull=max_hulls,
+        )
+
+        # Convert to trimesh objects
+        hulls = []
+        for verts, tris in parts:
+            hull = trimesh.Trimesh(vertices=verts, faces=tris)
+            hull.fix_normals()
+            hulls.append(hull)
+
+        logger.debug(f"CoACD decomposed into {len(hulls)} convex hulls")
+        return hulls
+
+    except ImportError:
+        logger.warning("CoACD not available, falling back to V-HACD")
+        return vhacd_decompose(mesh, max_hulls)
+
+
+def vhacd_decompose(
+    mesh: "trimesh.Trimesh",
+    max_hulls: int = 8,
+    resolution: int = 100000,
+) -> List["trimesh.Trimesh"]:
+    """
+    Decompose mesh into convex hulls using V-HACD.
+
+    V-HACD (Volumetric Hierarchical Approximate Convex Decomposition)
+    is a classic algorithm available in trimesh.
+
+    Args:
+        mesh: Input mesh to decompose
+        max_hulls: Maximum number of convex hulls
+        resolution: Voxelization resolution
+
+    Returns:
+        List of convex hull meshes
+    """
+    if not TRIMESH_AVAILABLE:
+        raise CollisionError("trimesh not available")
+
+    try:
+        # Try trimesh's built-in convex decomposition
+        hulls = mesh.convex_decomposition(
+            maxhulls=max_hulls,
+            resolution=resolution,
+        )
+
+        if isinstance(hulls, trimesh.Trimesh):
+            hulls = [hulls]
+
+        logger.debug(f"V-HACD decomposed into {len(hulls)} convex hulls")
+        return hulls
+
+    except Exception as e:
+        logger.warning(f"V-HACD failed: {e}, using single convex hull")
+        return [mesh.convex_hull]
+
+
+def simple_convex_hull(mesh: "trimesh.Trimesh") -> List["trimesh.Trimesh"]:
+    """
+    Create a single convex hull from mesh.
+
+    Simplest collision shape, fastest simulation.
+
+    Args:
+        mesh: Input mesh
+
+    Returns:
+        List containing single convex hull
+    """
+    if not TRIMESH_AVAILABLE:
+        raise CollisionError("trimesh not available")
+
+    hull = mesh.convex_hull
+    return [hull]
+
+
+def generate_collision_mesh(
+    visual_mesh_path: Path,
+    output_path: Path,
+    max_hulls: int = 8,
+    method: str = "auto",
+) -> Path:
+    """
+    Generate collision mesh from visual mesh.
+
+    Args:
+        visual_mesh_path: Path to visual mesh (OBJ)
+        output_path: Path for collision mesh output
+        max_hulls: Maximum convex hulls
+        method: "coacd", "vhacd", "hull", or "auto"
+
+    Returns:
+        Path to collision mesh
+    """
+    if not TRIMESH_AVAILABLE:
+        raise CollisionError("trimesh not available. Run: pip install trimesh")
+
+    # Load visual mesh
+    mesh = trimesh.load(str(visual_mesh_path))
+
+    if not isinstance(mesh, trimesh.Trimesh):
+        # Handle scene objects
+        if hasattr(mesh, 'geometry'):
+            meshes = list(mesh.geometry.values())
+            if meshes:
+                mesh = meshes[0]
+            else:
+                raise CollisionError(f"No geometry in {visual_mesh_path}")
+        else:
+            raise CollisionError(f"Invalid mesh type in {visual_mesh_path}")
+
+    # Decompose into convex hulls
+    if method == "auto":
+        try:
+            hulls = coacd_decompose(mesh, max_hulls)
+        except Exception:
+            try:
+                hulls = vhacd_decompose(mesh, max_hulls)
+            except Exception:
+                hulls = simple_convex_hull(mesh)
+    elif method == "coacd":
+        hulls = coacd_decompose(mesh, max_hulls)
+    elif method == "vhacd":
+        hulls = vhacd_decompose(mesh, max_hulls)
+    elif method == "hull":
+        hulls = simple_convex_hull(mesh)
+    else:
+        raise CollisionError(f"Unknown decomposition method: {method}")
+
+    # Combine hulls into single mesh for export
+    if len(hulls) == 1:
+        combined = hulls[0]
+    else:
+        combined = trimesh.util.concatenate(hulls)
+
+    # Export
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    combined.export(str(output_path))
+
+    logger.info(
+        f"Generated collision mesh: {output_path} "
+        f"({len(hulls)} hulls, {len(combined.faces)} faces)"
+    )
+    return output_path
+
+
+def generate_all_collision_meshes(
+    session: "ScanSession",
+    max_hulls: int = 8,
+    method: str = "auto",
+    progress_callback: Optional[callable] = None,
+) -> Dict[str, Path]:
+    """
+    Generate collision meshes for all links.
+
+    Args:
+        session: ScanSession with visual meshes
+        max_hulls: Maximum hulls per link
+        method: Decomposition method
+        progress_callback: Optional progress callback
+
+    Returns:
+        Dict mapping link_name -> collision_mesh_path
+    """
+    meshes_dir = session.meshes_dir
+    visual_meshes = list(meshes_dir.glob("*_visual.obj"))
+
+    if not visual_meshes:
+        raise CollisionError(f"No visual meshes found in {meshes_dir}")
+
+    result = {}
+    total = len(visual_meshes)
+
+    for i, visual_path in enumerate(visual_meshes):
+        link_name = visual_path.stem.replace("_visual", "")
+        collision_path = meshes_dir / f"{link_name}_collision.obj"
+
+        try:
+            generate_collision_mesh(visual_path, collision_path, max_hulls, method)
+            result[link_name] = collision_path
+        except Exception as e:
+            logger.error(f"Failed to generate collision mesh for {link_name}: {e}")
+
+        if progress_callback:
+            progress_callback(i + 1, total)
+
+    # Update session
+    session.metadata.mesh_complete = True
+    session.save_metadata()
+
+    logger.info(f"Generated {len(result)} collision meshes")
+    return result
+
+
+__all__ = [
+    "CollisionError",
+    "coacd_decompose",
+    "vhacd_decompose",
+    "simple_convex_hull",
+    "generate_collision_mesh",
+    "generate_all_collision_meshes",
+]