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

ate/urdf/segmentation.py

@@ -0,0 +1,299 @@
+"""
+Temporal segmentation for URDF generation using SAM 2.
+
+This module handles Phase 2a of the pipeline:
+1. Load SAM 2 model (with caching)
+2. Initialize masks from user click points
+3. Propagate masks across all video frames
+
+Output: Per-frame binary masks for each annotated link.
+"""
+
+import logging
+from typing import Dict, 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 SegmentationError(Exception):
+    """Error during segmentation."""
+    pass
+
+
+class LinkSegmenter:
+    """
+    Temporal segmentation of robot links using SAM 2.
+
+    Given a video and initial click points, propagates segmentation
+    masks for each link across all frames.
+    """
+
+    def __init__(self, device: str = "cpu"):
+        """
+        Initialize segmenter.
+
+        Args:
+            device: Compute device ("cuda" or "cpu")
+        """
+        if not DEPS_AVAILABLE:
+            raise SegmentationError(
+                "Required dependencies not installed. "
+                "Run: pip install numpy opencv-python"
+            )
+
+        self.device = device
+        self._predictor = None
+
+    def _get_predictor(self):
+        """Get or load the SAM 2 predictor."""
+        if self._predictor is None:
+            from .models import ModelCache
+            ModelCache.set_device(self.device)
+            self._predictor = ModelCache.get_sam2()
+        return self._predictor
+
+    def segment_video(
+        self,
+        video_path: str,
+        links: List[Tuple[str, List[float], bool]],
+        progress_callback: Optional[callable] = None,
+    ) -> Dict[str, Dict[int, "np.ndarray"]]:
+        """
+        Segment a video given initial link annotations.
+
+        Args:
+            video_path: Path to video file
+            links: List of (name, [x, y], is_fixed) tuples
+            progress_callback: Optional callback(current, total)
+
+        Returns:
+            Dict mapping link_name -> {frame_idx: mask}
+        """
+        predictor = self._get_predictor()
+
+        # Initialize predictor with video
+        predictor.initialize(video_path)
+
+        # Add point prompts for each link
+        logger.info(f"Initializing {len(links)} link masks...")
+        for obj_id, (name, point, is_fixed) in enumerate(links):
+            predictor.add_point_prompt(
+                frame_idx=0,
+                obj_id=obj_id,
+                point=(point[0], point[1]),
+                label=1,  # Foreground
+            )
+            logger.debug(f"Added prompt for '{name}' at {point}")
+
+        # Propagate masks
+        logger.info("Propagating masks across video frames...")
+        all_masks = predictor.propagate()
+
+        # Reorganize by link name
+        link_names = [name for name, _, _ in links]
+        result = {name: {} for name in link_names}
+
+        for frame_idx, frame_masks in all_masks.items():
+            for obj_id, mask in frame_masks.items():
+                if obj_id < len(link_names):
+                    link_name = link_names[obj_id]
+                    result[link_name][frame_idx] = mask
+
+            if progress_callback:
+                progress_callback(frame_idx + 1, len(all_masks))
+
+        # Validate results
+        for name in link_names:
+            if not result[name]:
+                raise SegmentationError(
+                    f"No masks generated for link '{name}'. "
+                    "Check that the click point is on the robot."
+                )
+
+        logger.info(f"Segmentation complete: {len(all_masks)} frames")
+        return result
+
+    def extract_masked_regions(
+        self,
+        video_path: str,
+        masks: Dict[str, Dict[int, "np.ndarray"]],
+        output_dir: Path,
+    ) -> Dict[str, List[Path]]:
+        """
+        Extract masked RGB regions from video frames.
+
+        Args:
+            video_path: Path to video file
+            masks: Dict from segment_video()
+            output_dir: Directory to save extracted regions
+
+        Returns:
+            Dict mapping link_name -> list of image paths
+        """
+        output_dir = Path(output_dir)
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        cap = cv2.VideoCapture(video_path)
+        if not cap.isOpened():
+            raise SegmentationError(f"Could not open video: {video_path}")
+
+        result = {name: [] for name in masks.keys()}
+        frame_idx = 0
+
+        try:
+            while True:
+                ret, frame = cap.read()
+                if not ret:
+                    break
+
+                for link_name, link_masks in masks.items():
+                    if frame_idx in link_masks:
+                        mask = link_masks[frame_idx]
+
+                        # Apply mask to frame
+                        masked = frame.copy()
+                        masked[~mask] = 0
+
+                        # Crop to bounding box
+                        coords = np.where(mask)
+                        if len(coords[0]) > 0:
+                            y_min, y_max = coords[0].min(), coords[0].max()
+                            x_min, x_max = coords[1].min(), coords[1].max()
+                            cropped = masked[y_min:y_max, x_min:x_max]
+
+                            # Save
+                            out_path = output_dir / f"{link_name}_frame_{frame_idx:04d}.png"
+                            cv2.imwrite(str(out_path), cropped)
+                            result[link_name].append(out_path)
+
+                frame_idx += 1
+
+        finally:
+            cap.release()
+
+        logger.info(f"Extracted masked regions to {output_dir}")
+        return result
+
+
+def visualize_masks(
+    frame: "np.ndarray",
+    masks: Dict[str, "np.ndarray"],
+    alpha: float = 0.5,
+) -> "np.ndarray":
+    """
+    Visualize segmentation masks overlaid on a frame.
+
+    Args:
+        frame: BGR image
+        masks: Dict mapping link_name -> mask
+        alpha: Overlay transparency
+
+    Returns:
+        Visualization image
+    """
+    if not DEPS_AVAILABLE:
+        raise SegmentationError("OpenCV not available")
+
+    display = frame.copy()
+
+    # Color palette for different links
+    colors = [
+        (255, 0, 0),    # Blue
+        (0, 255, 0),    # Green
+        (0, 0, 255),    # Red
+        (255, 255, 0),  # Cyan
+        (255, 0, 255),  # Magenta
+        (0, 255, 255),  # Yellow
+        (128, 0, 128),  # Purple
+        (255, 165, 0),  # Orange
+    ]
+
+    for i, (name, mask) in enumerate(masks.items()):
+        color = colors[i % len(colors)]
+
+        # Create colored overlay
+        overlay = display.copy()
+        overlay[mask] = color
+
+        # Blend
+        display = cv2.addWeighted(overlay, alpha, display, 1 - alpha, 0)
+
+        # Draw contour
+        contours, _ = cv2.findContours(
+            mask.astype(np.uint8),
+            cv2.RETR_EXTERNAL,
+            cv2.CHAIN_APPROX_SIMPLE,
+        )
+        cv2.drawContours(display, contours, -1, color, 2)
+
+        # Label
+        if contours:
+            M = cv2.moments(contours[0])
+            if M["m00"] > 0:
+                cx = int(M["m10"] / M["m00"])
+                cy = int(M["m01"] / M["m00"])
+                cv2.putText(
+                    display, name, (cx - 30, cy),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2
+                )
+                cv2.putText(
+                    display, name, (cx - 30, cy),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 1
+                )
+
+    return display
+
+
+def run_segmentation(
+    session: "ScanSession",
+    progress_callback: Optional[callable] = None,
+) -> Dict[str, Dict[int, "np.ndarray"]]:
+    """
+    Run segmentation phase on a scan session.
+
+    Args:
+        session: ScanSession with captured video and link annotations
+        progress_callback: Optional progress callback
+
+    Returns:
+        Dict mapping link_name -> {frame_idx: mask}
+    """
+    from .scan_session import ScanSession
+
+    # Check prerequisites
+    session.check_prerequisites("segment")
+
+    # Prepare link data
+    links = [
+        (link.name, link.point, link.is_fixed)
+        for link in session.links
+    ]
+
+    # Run segmentation
+    segmenter = LinkSegmenter(device=session.metadata.device)
+    masks = segmenter.segment_video(
+        str(session.video_path),
+        links,
+        progress_callback,
+    )
+
+    logger.info(f"Segmentation complete for {len(masks)} links")
+    return masks
+
+
+__all__ = [
+    "SegmentationError",
+    "LinkSegmenter",
+    "visualize_masks",
+    "run_segmentation",
+]

ate/urdf/synthesis.py

@@ -0,0 +1,319 @@
+"""
+URDF synthesis from discovered kinematics and generated meshes.
+
+This module handles Phase 5b of the pipeline:
+1. Load kinematics, meshes, and inertial properties
+2. Generate URDF XML structure
+3. Validate and write final file
+
+Output: Complete URDF file ready for simulation.
+"""
+
+import logging
+from typing import Dict, List, Optional
+from pathlib import Path
+from datetime import datetime
+import xml.etree.ElementTree as ET
+from xml.dom import minidom
+
+logger = logging.getLogger(__name__)
+
+
+class SynthesisError(Exception):
+    """Error during URDF synthesis."""
+    pass
+
+
+def create_urdf_link(
+    name: str,
+    visual_mesh: Optional[Path] = None,
+    collision_mesh: Optional[Path] = None,
+    inertial: Optional["InertialProperties"] = None,
+    mesh_base_path: str = "meshes",
+) -> ET.Element:
+    """
+    Create a URDF <link> element.
+
+    Args:
+        name: Link name
+        visual_mesh: Path to visual mesh file
+        collision_mesh: Path to collision mesh file
+        inertial: Inertial properties
+        mesh_base_path: Relative path prefix for mesh files
+
+    Returns:
+        ElementTree Element for the link
+    """
+    link = ET.Element("link", name=name)
+
+    # Visual geometry
+    if visual_mesh:
+        visual = ET.SubElement(link, "visual")
+        geometry = ET.SubElement(visual, "geometry")
+        mesh_path = f"{mesh_base_path}/{visual_mesh.name}"
+        ET.SubElement(geometry, "mesh", filename=mesh_path)
+
+        # Default material
+        material = ET.SubElement(visual, "material", name=f"{name}_material")
+        ET.SubElement(material, "color", rgba="0.7 0.7 0.7 1.0")
+
+    # Collision geometry
+    if collision_mesh:
+        collision = ET.SubElement(link, "collision")
+        geometry = ET.SubElement(collision, "geometry")
+        mesh_path = f"{mesh_base_path}/{collision_mesh.name}"
+        ET.SubElement(geometry, "mesh", filename=mesh_path)
+
+    # Inertial properties
+    if inertial:
+        inertial_elem = ET.SubElement(link, "inertial")
+        ET.SubElement(inertial_elem, "mass", value=f"{inertial.mass:.6f}")
+
+        com = inertial.center_of_mass
+        ET.SubElement(
+            inertial_elem, "origin",
+            xyz=f"{com[0]:.6f} {com[1]:.6f} {com[2]:.6f}",
+            rpy="0 0 0",
+        )
+
+        inertia = inertial.inertia
+        ET.SubElement(
+            inertial_elem, "inertia",
+            ixx=f"{inertia['ixx']:.9f}",
+            ixy=f"{inertia['ixy']:.9f}",
+            ixz=f"{inertia['ixz']:.9f}",
+            iyy=f"{inertia['iyy']:.9f}",
+            iyz=f"{inertia['iyz']:.9f}",
+            izz=f"{inertia['izz']:.9f}",
+        )
+
+    return link
+
+
+def create_urdf_joint(
+    name: str,
+    parent_link: str,
+    child_link: str,
+    joint_type: str,
+    axis: List[float],
+    origin: List[float],
+    limits: Dict[str, float],
+    effort: float = 100.0,
+    velocity: float = 1.0,
+) -> ET.Element:
+    """
+    Create a URDF <joint> element.
+
+    Args:
+        name: Joint name
+        parent_link: Parent link name
+        child_link: Child link name
+        joint_type: Joint type (revolute, prismatic, fixed)
+        axis: Rotation/translation axis [x, y, z]
+        origin: Joint origin [x, y, z]
+        limits: Joint limits {lower, upper}
+        effort: Maximum effort (N or Nm)
+        velocity: Maximum velocity (rad/s or m/s)
+
+    Returns:
+        ElementTree Element for the joint
+    """
+    joint = ET.Element("joint", name=name, type=joint_type)
+
+    ET.SubElement(joint, "parent", link=parent_link)
+    ET.SubElement(joint, "child", link=child_link)
+
+    ET.SubElement(
+        joint, "origin",
+        xyz=f"{origin[0]:.6f} {origin[1]:.6f} {origin[2]:.6f}",
+        rpy="0 0 0",
+    )
+
+    if joint_type != "fixed":
+        ET.SubElement(
+            joint, "axis",
+            xyz=f"{axis[0]:.6f} {axis[1]:.6f} {axis[2]:.6f}",
+        )
+
+        ET.SubElement(
+            joint, "limit",
+            lower=f"{limits['lower']:.6f}",
+            upper=f"{limits['upper']:.6f}",
+            effort=f"{effort:.1f}",
+            velocity=f"{velocity:.2f}",
+        )
+
+    return joint
+
+
+def synthesize_urdf(
+    session: "ScanSession",
+    inertials: Dict[str, "InertialProperties"],
+    robot_name: Optional[str] = None,
+    mesh_base_path: str = "meshes",
+) -> str:
+    """
+    Synthesize complete URDF from session data.
+
+    Args:
+        session: ScanSession with kinematics
+        inertials: Dict of InertialProperties per link
+        robot_name: Name for the robot (uses session name if not provided)
+        mesh_base_path: Relative path prefix for mesh files
+
+    Returns:
+        URDF XML string
+    """
+    robot_name = robot_name or session.metadata.robot_name or "robot"
+
+    # Create root element
+    robot = ET.Element("robot", name=robot_name)
+
+    # Add comment with generation info
+    comment = ET.Comment(
+        f"\n    Generated by ate urdf scan\n"
+        f"    Source: {session.session_dir}\n"
+        f"    Date: {datetime.now().isoformat()}\n  "
+    )
+    robot.insert(0, comment)
+
+    # Get mesh paths
+    meshes_dir = session.meshes_dir
+
+    # Create links
+    link_names = set()
+    for link in session.links:
+        link_names.add(link.name)
+    for joint in session.joints:
+        link_names.add(joint.parent_link)
+        link_names.add(joint.child_link)
+
+    for link_name in sorted(link_names):
+        visual_mesh = meshes_dir / f"{link_name}_visual.obj"
+        collision_mesh = meshes_dir / f"{link_name}_collision.obj"
+
+        visual_path = visual_mesh if visual_mesh.exists() else None
+        collision_path = collision_mesh if collision_mesh.exists() else None
+        inertial = inertials.get(link_name)
+
+        link_elem = create_urdf_link(
+            link_name,
+            visual_path,
+            collision_path,
+            inertial,
+            mesh_base_path,
+        )
+        robot.append(link_elem)
+
+    # Create joints
+    for joint in session.joints:
+        joint_elem = create_urdf_joint(
+            name=joint.name,
+            parent_link=joint.parent_link,
+            child_link=joint.child_link,
+            joint_type=joint.joint_type,
+            axis=joint.axis,
+            origin=joint.origin,
+            limits=joint.limits,
+        )
+        robot.append(joint_elem)
+
+    # Convert to string with pretty formatting
+    xml_string = ET.tostring(robot, encoding='unicode')
+    dom = minidom.parseString(xml_string)
+    pretty_xml = dom.toprettyxml(indent="  ")
+
+    # Remove extra blank lines
+    lines = [line for line in pretty_xml.split('\n') if line.strip()]
+    # Re-add XML declaration
+    pretty_xml = '\n'.join(lines)
+
+    return pretty_xml
+
+
+def write_urdf(
+    session: "ScanSession",
+    inertials: Dict[str, "InertialProperties"],
+    output_path: Optional[Path] = None,
+    robot_name: Optional[str] = None,
+) -> Path:
+    """
+    Write URDF file to disk.
+
+    Args:
+        session: ScanSession
+        inertials: Inertial properties
+        output_path: Output path (uses session default if not provided)
+        robot_name: Robot name
+
+    Returns:
+        Path to written URDF file
+    """
+    output_path = output_path or session.urdf_path
+
+    urdf_content = synthesize_urdf(session, inertials, robot_name)
+
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(output_path, 'w') as f:
+        f.write(urdf_content)
+
+    logger.info(f"Wrote URDF to {output_path}")
+    return output_path
+
+
+def run_synthesis(
+    session: "ScanSession",
+    density: float = 1200.0,
+    robot_name: Optional[str] = None,
+    progress_callback: Optional[callable] = None,
+) -> Path:
+    """
+    Run the URDF synthesis phase.
+
+    Args:
+        session: ScanSession with kinematics and meshes
+        density: Material density for inertial estimation
+        robot_name: Robot name
+        progress_callback: Optional progress callback
+
+    Returns:
+        Path to generated URDF
+    """
+    from .inertial import estimate_all_inertials
+
+    # Check prerequisites
+    session.check_prerequisites("synthesize")
+
+    # Estimate inertials
+    logger.info("Estimating inertial properties...")
+    if progress_callback:
+        progress_callback(1, 3)
+
+    inertials = estimate_all_inertials(session, density)
+
+    # Synthesize URDF
+    logger.info("Synthesizing URDF...")
+    if progress_callback:
+        progress_callback(2, 3)
+
+    urdf_path = write_urdf(session, inertials, robot_name=robot_name)
+
+    # Update session
+    session.metadata.synthesize_complete = True
+    session.save_metadata()
+
+    if progress_callback:
+        progress_callback(3, 3)
+
+    logger.info(f"URDF synthesis complete: {urdf_path}")
+    return urdf_path
+
+
+__all__ = [
+    "SynthesisError",
+    "create_urdf_link",
+    "create_urdf_joint",
+    "synthesize_urdf",
+    "write_urdf",
+    "run_synthesis",
+]