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

ate/urdf/scale.py

@@ -0,0 +1,256 @@
+"""
+Scale calibration for markerless URDF generation.
+
+In a markerless setting, the global scale factor (sim2real gap) must be
+determined from a known physical measurement. This module provides:
+
+- parse_scale_ref: Parse "part:measurement" strings (e.g., "gripper:85mm")
+- parse_measurement: Parse measurement values with units
+- ScaleCalibrator: Compute scale factor from reference measurements
+
+The Kalib method uses a single known dimension to normalize depth estimates
+from monocular depth models like Depth Anything V2.
+"""
+
+import re
+from typing import Tuple, Optional
+from dataclasses import dataclass
+
+
+class ScaleError(Exception):
+    """Error in scale parsing or calibration."""
+    pass
+
+
+# Unit conversion factors to meters
+UNIT_CONVERSIONS = {
+    # Metric
+    "m": 1.0,
+    "meter": 1.0,
+    "meters": 1.0,
+    "cm": 0.01,
+    "centimeter": 0.01,
+    "centimeters": 0.01,
+    "mm": 0.001,
+    "millimeter": 0.001,
+    "millimeters": 0.001,
+    # Imperial
+    "in": 0.0254,
+    "inch": 0.0254,
+    "inches": 0.0254,
+    '"': 0.0254,
+    "ft": 0.3048,
+    "foot": 0.3048,
+    "feet": 0.3048,
+    "'": 0.3048,
+}
+
+
+def parse_measurement(measurement: str) -> float:
+    """
+    Parse a measurement string with units to meters.
+
+    Args:
+        measurement: Value with unit (e.g., "85mm", "3.5in", "0.15m")
+
+    Returns:
+        Value in meters
+
+    Raises:
+        ScaleError: If format is invalid or unit is unknown
+
+    Examples:
+        >>> parse_measurement("85mm")
+        0.085
+        >>> parse_measurement("3.5in")
+        0.0889
+        >>> parse_measurement("150cm")
+        1.5
+    """
+    if not measurement:
+        raise ScaleError("Empty measurement string")
+
+    # Normalize: strip whitespace, lowercase
+    measurement = measurement.strip().lower()
+
+    # Pattern: number (with optional decimal) followed by unit
+    pattern = r"^([\d.]+)\s*([a-z\"']+)$"
+    match = re.match(pattern, measurement)
+
+    if not match:
+        raise ScaleError(
+            f"Invalid measurement format: '{measurement}'. "
+            f"Expected format like '85mm', '3.5in', '0.15m'"
+        )
+
+    value_str, unit = match.groups()
+
+    try:
+        value = float(value_str)
+    except ValueError:
+        raise ScaleError(f"Invalid numeric value: '{value_str}'")
+
+    if value <= 0:
+        raise ScaleError(f"Measurement must be positive: {value}")
+
+    if unit not in UNIT_CONVERSIONS:
+        valid_units = sorted(set(UNIT_CONVERSIONS.keys()) - {'"', "'"})
+        raise ScaleError(
+            f"Unknown unit: '{unit}'. Valid units: {', '.join(valid_units)}"
+        )
+
+    return value * UNIT_CONVERSIONS[unit]
+
+
+def parse_scale_ref(ref: str) -> Tuple[str, float]:
+    """
+    Parse a scale reference string.
+
+    Format: "part_name:measurement"
+
+    Args:
+        ref: Scale reference (e.g., "gripper:85mm", "base_width:150mm")
+
+    Returns:
+        Tuple of (part_name, value_in_meters)
+
+    Raises:
+        ScaleError: If format is invalid
+
+    Examples:
+        >>> parse_scale_ref("gripper:85mm")
+        ('gripper', 0.085)
+        >>> parse_scale_ref("base_width:6in")
+        ('base_width', 0.1524)
+    """
+    if not ref:
+        raise ScaleError("Empty scale reference")
+
+    if ":" not in ref:
+        raise ScaleError(
+            f"Invalid scale reference format: '{ref}'. "
+            f"Expected 'part:measurement' (e.g., 'gripper:85mm')"
+        )
+
+    parts = ref.split(":", 1)
+    if len(parts) != 2:
+        raise ScaleError(f"Invalid scale reference: '{ref}'")
+
+    part_name = parts[0].strip()
+    measurement = parts[1].strip()
+
+    if not part_name:
+        raise ScaleError("Part name cannot be empty in scale reference")
+
+    value_meters = parse_measurement(measurement)
+
+    return part_name, value_meters
+
+
+@dataclass
+class ScaleCalibration:
+    """Result of scale calibration."""
+    part_name: str
+    reference_meters: float
+    measured_pixels: float
+    scale_factor: float  # meters per pixel (or depth unit)
+
+    def apply(self, depth_value: float) -> float:
+        """Apply scale factor to convert raw depth to meters."""
+        return depth_value * self.scale_factor
+
+
+class ScaleCalibrator:
+    """
+    Calibrate global scale from a known physical measurement.
+
+    Uses the Kalib method: correlate a known physical dimension with
+    its measured size in the depth map or point cloud.
+
+    Usage:
+        calibrator = ScaleCalibrator("gripper:85mm")
+        # Measure gripper width in depth units
+        measured = compute_dimension_from_cloud(gripper_cloud)
+        calibration = calibrator.calibrate(measured)
+        # Apply to all depth values
+        scaled_depth = calibration.apply(raw_depth)
+    """
+
+    def __init__(self, scale_ref: str):
+        """
+        Initialize calibrator with a scale reference.
+
+        Args:
+            scale_ref: Scale reference string (e.g., "gripper:85mm")
+        """
+        self.part_name, self.reference_meters = parse_scale_ref(scale_ref)
+
+    def calibrate(self, measured_value: float) -> ScaleCalibration:
+        """
+        Compute scale factor from measured value.
+
+        Args:
+            measured_value: Size of reference part in raw depth units
+
+        Returns:
+            ScaleCalibration with computed scale factor
+        """
+        if measured_value <= 0:
+            raise ScaleError(
+                f"Measured value must be positive: {measured_value}. "
+                f"Check that the reference part '{self.part_name}' is visible."
+            )
+
+        scale_factor = self.reference_meters / measured_value
+
+        return ScaleCalibration(
+            part_name=self.part_name,
+            reference_meters=self.reference_meters,
+            measured_pixels=measured_value,
+            scale_factor=scale_factor,
+        )
+
+
+def estimate_intrinsics_from_resolution(
+    width: int,
+    height: int,
+    fov_degrees: float = 60.0,
+) -> Tuple[float, float, float, float]:
+    """
+    Estimate camera intrinsics from resolution and assumed FOV.
+
+    This is a fallback when actual camera calibration is not available.
+    Most webcams have FOV between 55-75 degrees.
+
+    Args:
+        width: Image width in pixels
+        height: Image height in pixels
+        fov_degrees: Assumed horizontal field of view
+
+    Returns:
+        Tuple of (fx, fy, cx, cy) camera intrinsics
+    """
+    import math
+
+    # Compute focal length from FOV
+    fov_rad = math.radians(fov_degrees)
+    fx = (width / 2) / math.tan(fov_rad / 2)
+
+    # Assume square pixels
+    fy = fx
+
+    # Principal point at image center
+    cx = width / 2
+    cy = height / 2
+
+    return fx, fy, cx, cy
+
+
+__all__ = [
+    "ScaleError",
+    "parse_measurement",
+    "parse_scale_ref",
+    "ScaleCalibration",
+    "ScaleCalibrator",
+    "estimate_intrinsics_from_resolution",
+]

ate/urdf/scan_session.py

@@ -0,0 +1,411 @@
+"""
+Scan session management for URDF generation.
+
+A scan session represents the state of a URDF generation process,
+including captured video, link annotations, point clouds, meshes,
+and final URDF output.
+
+Session structure:
+    my_robot_scan/
+    ├── metadata.json        # Session metadata and settings
+    ├── video.mp4           # Captured robot video
+    ├── links.json          # User-annotated link click points
+    ├── kinematics.json     # Discovered joint parameters
+    ├── clouds/             # Per-link point clouds
+    │   ├── base_frame_0.ply
+    │   ├── shoulder_frame_0.ply
+    │   └── ...
+    └── meshes/             # Generated meshes
+        ├── base_visual.obj
+        ├── base_collision.obj
+        └── ...
+"""
+
+import json
+import os
+from dataclasses import dataclass, field, asdict
+from datetime import datetime
+from pathlib import Path
+from typing import Dict, List, Optional, Any
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class ScanSessionError(Exception):
+    """Base exception for scan session errors."""
+    pass
+
+
+class SessionNotFoundError(ScanSessionError):
+    """Session directory does not exist."""
+    pass
+
+
+class SessionCorruptError(ScanSessionError):
+    """Session metadata is missing or invalid."""
+    pass
+
+
+class SessionIncompleteError(ScanSessionError):
+    """Session is missing required data for the requested operation."""
+
+    def __init__(self, message: str, missing_steps: List[str]):
+        super().__init__(message)
+        self.missing_steps = missing_steps
+
+
+@dataclass
+class LinkAnnotation:
+    """A user-annotated link in the first video frame."""
+    name: str
+    point: List[float]  # [x, y] coordinates in frame
+    is_fixed: bool = False  # True for base/world link
+
+    def to_dict(self) -> Dict:
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: Dict) -> "LinkAnnotation":
+        return cls(
+            name=data["name"],
+            point=data["point"],
+            is_fixed=data.get("is_fixed", False),
+        )
+
+
+@dataclass
+class JointInfo:
+    """Discovered joint parameters."""
+    name: str
+    parent_link: str
+    child_link: str
+    joint_type: str  # revolute, prismatic, fixed
+    axis: List[float]  # [x, y, z] unit vector
+    origin: List[float]  # [x, y, z] position
+    limits: Dict[str, float]  # lower, upper bounds (radians or meters)
+
+    def to_dict(self) -> Dict:
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: Dict) -> "JointInfo":
+        return cls(
+            name=data["name"],
+            parent_link=data["parent_link"],
+            child_link=data["child_link"],
+            joint_type=data["type"],
+            axis=data["axis"],
+            origin=data["origin"],
+            limits=data.get("limits", {"lower": -3.14, "upper": 3.14}),
+        )
+
+
+@dataclass
+class ScanMetadata:
+    """Session metadata and settings."""
+    version: str = "1.0.0"
+    created_at: str = field(default_factory=lambda: datetime.now().isoformat())
+    updated_at: str = field(default_factory=lambda: datetime.now().isoformat())
+    robot_name: Optional[str] = None
+    scale_ref: Optional[str] = None  # e.g., "gripper:85mm"
+    scale_factor: Optional[float] = None  # Computed from scale_ref
+    device: str = "cpu"  # cpu or cuda
+    video_path: Optional[str] = None
+    frame_count: int = 0
+    fps: float = 30.0
+    resolution: List[int] = field(default_factory=lambda: [0, 0])  # [width, height]
+    density_kg_m3: float = 1200.0  # Default plastic density
+
+    # Pipeline completion status
+    capture_complete: bool = False
+    segment_complete: bool = False
+    optimize_complete: bool = False
+    mesh_complete: bool = False
+    synthesize_complete: bool = False
+
+    def to_dict(self) -> Dict:
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: Dict) -> "ScanMetadata":
+        return cls(**{k: v for k, v in data.items() if k in cls.__dataclass_fields__})
+
+    def update(self) -> None:
+        """Update the modification timestamp."""
+        self.updated_at = datetime.now().isoformat()
+
+
+class ScanSession:
+    """
+    Manages a URDF generation session.
+
+    Provides methods for:
+    - Creating new sessions
+    - Loading existing sessions
+    - Saving/loading video, annotations, point clouds, meshes
+    - Tracking pipeline progress
+    """
+
+    METADATA_FILE = "metadata.json"
+    VIDEO_FILE = "video.mp4"
+    LINKS_FILE = "links.json"
+    KINEMATICS_FILE = "kinematics.json"
+    CLOUDS_DIR = "clouds"
+    MESHES_DIR = "meshes"
+    URDF_FILE = "robot.urdf"
+
+    def __init__(self, session_dir: Path):
+        """
+        Initialize a scan session.
+
+        Args:
+            session_dir: Directory for session data
+        """
+        self.session_dir = Path(session_dir)
+        self.metadata: ScanMetadata = ScanMetadata()
+        self.links: List[LinkAnnotation] = []
+        self.joints: List[JointInfo] = []
+
+    @classmethod
+    def create(
+        cls,
+        output_dir: str,
+        robot_name: Optional[str] = None,
+        scale_ref: Optional[str] = None,
+        device: str = "cpu",
+    ) -> "ScanSession":
+        """
+        Create a new scan session.
+
+        Args:
+            output_dir: Directory to create session in
+            robot_name: Name for the robot
+            scale_ref: Scale reference (e.g., "gripper:85mm")
+            device: Compute device (cpu or cuda)
+
+        Returns:
+            New ScanSession instance
+        """
+        session_dir = Path(output_dir)
+        session_dir.mkdir(parents=True, exist_ok=True)
+
+        # Create subdirectories
+        (session_dir / cls.CLOUDS_DIR).mkdir(exist_ok=True)
+        (session_dir / cls.MESHES_DIR).mkdir(exist_ok=True)
+
+        session = cls(session_dir)
+        session.metadata.robot_name = robot_name or session_dir.name
+        session.metadata.scale_ref = scale_ref
+        session.metadata.device = device
+
+        # Parse scale reference if provided
+        if scale_ref:
+            from .scale import parse_scale_ref
+            _, scale_meters = parse_scale_ref(scale_ref)
+            session.metadata.scale_factor = scale_meters
+
+        session.save_metadata()
+        logger.info(f"Created new scan session at {session_dir}")
+        return session
+
+    @classmethod
+    def load(cls, session_dir: str) -> "ScanSession":
+        """
+        Load an existing scan session.
+
+        Args:
+            session_dir: Path to session directory
+
+        Returns:
+            Loaded ScanSession instance
+
+        Raises:
+            SessionNotFoundError: If directory doesn't exist
+            SessionCorruptError: If metadata is invalid
+        """
+        path = Path(session_dir)
+
+        if not path.exists():
+            raise SessionNotFoundError(f"Session directory not found: {session_dir}")
+
+        metadata_path = path / cls.METADATA_FILE
+        if not metadata_path.exists():
+            raise SessionCorruptError(f"Missing metadata file: {metadata_path}")
+
+        session = cls(path)
+
+        # Load metadata
+        try:
+            with open(metadata_path, "r") as f:
+                data = json.load(f)
+                session.metadata = ScanMetadata.from_dict(data)
+        except json.JSONDecodeError as e:
+            raise SessionCorruptError(f"Invalid metadata JSON: {e}")
+
+        # Load links if present
+        links_path = path / cls.LINKS_FILE
+        if links_path.exists():
+            session.links = session._load_links()
+
+        # Load kinematics if present
+        kinematics_path = path / cls.KINEMATICS_FILE
+        if kinematics_path.exists():
+            session.joints = session._load_kinematics()
+
+        logger.info(f"Loaded scan session from {session_dir}")
+        return session
+
+    def save_metadata(self) -> None:
+        """Save session metadata to disk."""
+        self.metadata.update()
+        metadata_path = self.session_dir / self.METADATA_FILE
+        with open(metadata_path, "w") as f:
+            json.dump(self.metadata.to_dict(), f, indent=2)
+        logger.debug(f"Saved metadata to {metadata_path}")
+
+    def save_links(self) -> None:
+        """Save link annotations to disk."""
+        links_path = self.session_dir / self.LINKS_FILE
+        data = {
+            "frame_index": 0,
+            "links": [link.to_dict() for link in self.links],
+        }
+        with open(links_path, "w") as f:
+            json.dump(data, f, indent=2)
+        logger.debug(f"Saved {len(self.links)} links to {links_path}")
+
+    def _load_links(self) -> List[LinkAnnotation]:
+        """Load link annotations from disk."""
+        links_path = self.session_dir / self.LINKS_FILE
+        with open(links_path, "r") as f:
+            data = json.load(f)
+        return [LinkAnnotation.from_dict(link) for link in data.get("links", [])]
+
+    def save_kinematics(self) -> None:
+        """Save discovered kinematics to disk."""
+        kinematics_path = self.session_dir / self.KINEMATICS_FILE
+        data = {
+            "links": [
+                {"name": link.name, "parent": None if link.is_fixed else "discovered"}
+                for link in self.links
+            ],
+            "joints": [joint.to_dict() for joint in self.joints],
+        }
+        with open(kinematics_path, "w") as f:
+            json.dump(data, f, indent=2)
+        logger.debug(f"Saved {len(self.joints)} joints to {kinematics_path}")
+
+    def _load_kinematics(self) -> List[JointInfo]:
+        """Load kinematics from disk."""
+        kinematics_path = self.session_dir / self.KINEMATICS_FILE
+        with open(kinematics_path, "r") as f:
+            data = json.load(f)
+        return [JointInfo.from_dict(joint) for joint in data.get("joints", [])]
+
+    @property
+    def video_path(self) -> Path:
+        """Path to the video file."""
+        return self.session_dir / self.VIDEO_FILE
+
+    @property
+    def clouds_dir(self) -> Path:
+        """Path to point clouds directory."""
+        return self.session_dir / self.CLOUDS_DIR
+
+    @property
+    def meshes_dir(self) -> Path:
+        """Path to meshes directory."""
+        return self.session_dir / self.MESHES_DIR
+
+    @property
+    def urdf_path(self) -> Path:
+        """Path to final URDF file."""
+        return self.session_dir / self.URDF_FILE
+
+    def has_video(self) -> bool:
+        """Check if video has been captured."""
+        return self.video_path.exists()
+
+    def has_links(self) -> bool:
+        """Check if links have been annotated."""
+        return len(self.links) > 0
+
+    def has_clouds(self) -> bool:
+        """Check if point clouds have been generated."""
+        if not self.clouds_dir.exists():
+            return False
+        return len(list(self.clouds_dir.glob("*.ply"))) > 0
+
+    def has_meshes(self) -> bool:
+        """Check if meshes have been generated."""
+        if not self.meshes_dir.exists():
+            return False
+        return len(list(self.meshes_dir.glob("*_visual.obj"))) > 0
+
+    def has_urdf(self) -> bool:
+        """Check if URDF has been generated."""
+        return self.urdf_path.exists()
+
+    def check_prerequisites(self, stage: str) -> None:
+        """
+        Check if prerequisites for a pipeline stage are met.
+
+        Args:
+            stage: Pipeline stage to check
+
+        Raises:
+            SessionIncompleteError: If prerequisites are not met
+        """
+        missing = []
+
+        if stage == "segment":
+            if not self.has_video():
+                missing.append("video (run 'ate urdf scan capture' first)")
+            if not self.has_links():
+                missing.append("link annotations (annotate links during capture)")
+
+        elif stage == "optimize":
+            if not self.metadata.segment_complete:
+                missing.append("segmentation (run 'ate urdf scan segment' first)")
+
+        elif stage == "mesh":
+            if not self.metadata.optimize_complete:
+                missing.append("kinematics (run 'ate urdf scan optimize' first)")
+
+        elif stage == "synthesize":
+            if not self.metadata.mesh_complete:
+                missing.append("meshes (run 'ate urdf scan mesh' first)")
+
+        if missing:
+            raise SessionIncompleteError(
+                f"Cannot run '{stage}' stage. Missing: {', '.join(missing)}",
+                missing_steps=missing,
+            )
+
+    def get_status(self) -> Dict[str, Any]:
+        """Get current session status summary."""
+        return {
+            "session_dir": str(self.session_dir),
+            "robot_name": self.metadata.robot_name,
+            "scale_ref": self.metadata.scale_ref,
+            "device": self.metadata.device,
+            "stages": {
+                "capture": self.metadata.capture_complete,
+                "segment": self.metadata.segment_complete,
+                "optimize": self.metadata.optimize_complete,
+                "mesh": self.metadata.mesh_complete,
+                "synthesize": self.metadata.synthesize_complete,
+            },
+            "data": {
+                "has_video": self.has_video(),
+                "link_count": len(self.links),
+                "joint_count": len(self.joints),
+                "has_clouds": self.has_clouds(),
+                "has_meshes": self.has_meshes(),
+                "has_urdf": self.has_urdf(),
+            },
+        }
+
+    def __repr__(self) -> str:
+        return f"ScanSession({self.session_dir}, robot={self.metadata.robot_name})"