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

ate/urdf/pipeline.py

@@ -0,0 +1,468 @@
+"""
+Full URDF generation pipeline orchestration.
+
+This module provides the unified pipeline that runs all phases:
+1. Capture: Video capture and link annotation
+2. Segment: SAM 2 segmentation + Depth Anything depth estimation + 3D lifting
+3. Optimize: Kinematic parameter estimation
+4. Mesh: Visual and collision mesh generation
+5. Synthesize: URDF XML generation and validation
+
+Usage:
+    ate urdf scan --output ./my_robot/ --name my_robot --scale-ref "gripper:85mm"
+"""
+
+import logging
+import sys
+from typing import Optional, Dict, Any
+from pathlib import Path
+from dataclasses import dataclass
+
+logger = logging.getLogger(__name__)
+
+
+class PipelineError(Exception):
+    """Error during pipeline execution."""
+    pass
+
+
+@dataclass
+class PipelineConfig:
+    """Configuration for the URDF generation pipeline."""
+    output_dir: str
+    robot_name: Optional[str] = None
+    scale_ref: Optional[str] = None
+    video_path: Optional[str] = None
+    device: str = "cpu"
+    camera_id: int = 0
+    density: float = 1200.0  # kg/m^3
+    max_hulls: int = 8
+    frame_skip: int = 1
+    voxel_size: float = 0.005  # meters
+    simplify_to: int = 5000  # faces
+    dry_run: bool = False
+    upload: bool = False
+
+
+def create_progress_callback(stage_name: str):
+    """Create a progress callback for a pipeline stage."""
+    def callback(current: int, total: int):
+        pct = int(100 * current / total)
+        bar_len = 30
+        filled = int(bar_len * current / total)
+        bar = "=" * filled + "-" * (bar_len - filled)
+        print(f"\r  [{bar}] {pct}% ({current}/{total})", end="", flush=True)
+        if current >= total:
+            print()  # New line when complete
+    return callback
+
+
+def run_capture_stage(config: PipelineConfig) -> "ScanSession":
+    """
+    Run the capture stage.
+
+    Args:
+        config: Pipeline configuration
+
+    Returns:
+        ScanSession with captured data
+    """
+    from .capture import run_capture
+
+    print("\n" + "=" * 60)
+    print("PHASE 1: DATA ACQUISITION")
+    print("=" * 60)
+
+    if config.dry_run:
+        print("  [DRY RUN] Would capture video and annotate links")
+        from .scan_session import ScanSession
+        return ScanSession.create(
+            config.output_dir,
+            robot_name=config.robot_name,
+            scale_ref=config.scale_ref,
+            device=config.device,
+        )
+
+    session = run_capture(
+        output_dir=config.output_dir,
+        video_path=config.video_path,
+        camera_id=config.camera_id,
+        robot_name=config.robot_name,
+        scale_ref=config.scale_ref,
+    )
+
+    print(f"\nCapture complete: {len(session.links)} links annotated")
+    return session
+
+
+def run_segment_stage(
+    session: "ScanSession",
+    config: PipelineConfig,
+) -> Dict[str, Any]:
+    """
+    Run the segmentation and depth estimation stage.
+
+    Args:
+        session: ScanSession from capture stage
+        config: Pipeline configuration
+
+    Returns:
+        Dict with masks, depth_maps, and cloud_paths
+    """
+    from .segmentation import run_segmentation
+    from .depth import run_depth_estimation
+    from .lifting import run_lifting
+
+    print("\n" + "=" * 60)
+    print("PHASE 2: SEGMENTATION & DEPTH")
+    print("=" * 60)
+
+    if config.dry_run:
+        print("  [DRY RUN] Would run SAM 2 segmentation")
+        print("  [DRY RUN] Would run Depth Anything V2")
+        print("  [DRY RUN] Would generate point clouds")
+        return {"masks": {}, "depth_maps": {}, "cloud_paths": {}}
+
+    # Run segmentation
+    print("\nRunning SAM 2 segmentation...")
+    masks = run_segmentation(
+        session,
+        progress_callback=create_progress_callback("Segmentation"),
+    )
+
+    # Run depth estimation
+    print("\nRunning depth estimation...")
+    depth_maps, scale_factor = run_depth_estimation(
+        session,
+        masks,
+        frame_skip=config.frame_skip,
+        progress_callback=create_progress_callback("Depth"),
+    )
+
+    # Generate point clouds
+    print("\nGenerating point clouds...")
+    cloud_paths = run_lifting(
+        session,
+        masks,
+        depth_maps,
+        frame_skip=config.frame_skip,
+        voxel_size=config.voxel_size,
+        progress_callback=create_progress_callback("Lifting"),
+    )
+
+    print(f"\nSegmentation complete: {len(cloud_paths)} point clouds generated")
+    return {
+        "masks": masks,
+        "depth_maps": depth_maps,
+        "cloud_paths": cloud_paths,
+    }
+
+
+def run_optimize_stage(
+    session: "ScanSession",
+    segment_data: Dict[str, Any],
+    config: PipelineConfig,
+):
+    """
+    Run the kinematic optimization stage.
+
+    Args:
+        session: ScanSession
+        segment_data: Output from segment stage
+        config: Pipeline configuration
+    """
+    from .kinematics import run_kinematic_optimization
+
+    print("\n" + "=" * 60)
+    print("PHASE 3: KINEMATIC OPTIMIZATION")
+    print("=" * 60)
+
+    if config.dry_run:
+        print("  [DRY RUN] Would estimate joint parameters")
+        return
+
+    print("\nEstimating joint parameters...")
+    joints = run_kinematic_optimization(
+        session,
+        segment_data["masks"],
+        segment_data["depth_maps"],
+        progress_callback=create_progress_callback("Kinematics"),
+    )
+
+    print(f"\nOptimization complete: {len(joints)} joints discovered")
+    for joint in joints:
+        print(f"  - {joint.name}: {joint.joint_type} (conf={joint.confidence:.2f})")
+
+
+def run_mesh_stage(session: "ScanSession", config: PipelineConfig):
+    """
+    Run the mesh generation stage.
+
+    Args:
+        session: ScanSession
+        config: Pipeline configuration
+    """
+    from .meshing import generate_all_visual_meshes
+    from .collision import generate_all_collision_meshes
+
+    print("\n" + "=" * 60)
+    print("PHASE 4: MESH GENERATION")
+    print("=" * 60)
+
+    if config.dry_run:
+        print("  [DRY RUN] Would generate visual meshes")
+        print("  [DRY RUN] Would generate collision meshes")
+        return
+
+    print("\nGenerating visual meshes...")
+    visual_paths = generate_all_visual_meshes(
+        session,
+        simplify_to=config.simplify_to,
+        progress_callback=create_progress_callback("Visual meshes"),
+    )
+
+    print("\nGenerating collision meshes...")
+    collision_paths = generate_all_collision_meshes(
+        session,
+        max_hulls=config.max_hulls,
+        progress_callback=create_progress_callback("Collision meshes"),
+    )
+
+    print(f"\nMesh generation complete: {len(visual_paths)} links")
+
+
+def run_synthesize_stage(session: "ScanSession", config: PipelineConfig) -> Path:
+    """
+    Run the URDF synthesis stage.
+
+    Args:
+        session: ScanSession
+        config: Pipeline configuration
+
+    Returns:
+        Path to generated URDF
+    """
+    from .synthesis import run_synthesis
+    from .validation import run_validation
+
+    print("\n" + "=" * 60)
+    print("PHASE 5: URDF SYNTHESIS")
+    print("=" * 60)
+
+    if config.dry_run:
+        print("  [DRY RUN] Would synthesize URDF")
+        print("  [DRY RUN] Would validate URDF")
+        return session.urdf_path
+
+    print("\nSynthesizing URDF...")
+    urdf_path = run_synthesis(
+        session,
+        density=config.density,
+        robot_name=config.robot_name,
+        progress_callback=create_progress_callback("Synthesis"),
+    )
+
+    print("\nValidating URDF...")
+    result = run_validation(session)
+
+    if result.valid:
+        print("  URDF validation: PASSED")
+    else:
+        print("  URDF validation: FAILED")
+        for error in result.errors:
+            print(f"    ERROR: {error}")
+
+    for warning in result.warnings:
+        print(f"    WARNING: {warning}")
+
+    print(f"\nGenerated: {urdf_path}")
+    return urdf_path
+
+
+def run_full_pipeline(config: PipelineConfig) -> Path:
+    """
+    Run the complete URDF generation pipeline.
+
+    Args:
+        config: Pipeline configuration
+
+    Returns:
+        Path to generated URDF file
+    """
+    print("\n" + "#" * 60)
+    print("# MARKERLESS URDF GENERATION PIPELINE")
+    print("#" * 60)
+    print(f"\nOutput: {config.output_dir}")
+    print(f"Robot: {config.robot_name or 'auto'}")
+    print(f"Scale: {config.scale_ref or 'unspecified'}")
+    print(f"Device: {config.device}")
+
+    if config.dry_run:
+        print("\n*** DRY RUN MODE - No changes will be made ***")
+
+    try:
+        # Phase 1: Capture
+        session = run_capture_stage(config)
+
+        # Phase 2: Segment
+        segment_data = run_segment_stage(session, config)
+
+        # Phase 3: Optimize
+        run_optimize_stage(session, segment_data, config)
+
+        # Phase 4: Mesh
+        run_mesh_stage(session, config)
+
+        # Phase 5: Synthesize
+        urdf_path = run_synthesize_stage(session, config)
+
+        # Summary
+        print("\n" + "#" * 60)
+        print("# PIPELINE COMPLETE")
+        print("#" * 60)
+        status = session.get_status()
+        print(f"\nSession: {status['session_dir']}")
+        print(f"Links: {status['data']['link_count']}")
+        print(f"Joints: {status['data']['joint_count']}")
+        print(f"URDF: {urdf_path}")
+
+        if config.upload:
+            print("\nUploading to FoodforThought...")
+            try:
+                from ..robot.skill_upload import SkillLibraryUploader, APIError
+
+                uploader = SkillLibraryUploader()
+                robot_name = session.metadata.robot_name or "unknown_robot"
+                
+                project = uploader.get_or_create_project(
+                    name=f"{robot_name}_urdf",
+                    description=f"URDF from markerless scan for {robot_name}",
+                )
+                
+                urdf_content = urdf_path.read_text()
+                response = uploader._request(
+                    "POST",
+                    f"/projects/{project['id']}/artifacts",
+                    json={
+                        "name": f"{robot_name}.urdf",
+                        "artifact_type": "processed",
+                        "content_type": "application/xml",
+                        "metadata": {
+                            "robot_name": robot_name,
+                            "scale_ref": session.metadata.scale_ref,
+                            "link_count": len(session.links),
+                            "joint_count": len(session.joints),
+                            "generated_by": "ate urdf scan",
+                        },
+                        "content": urdf_content,
+                    },
+                )
+                print(f"Uploaded: {robot_name}.urdf (artifact {response.get('id', 'unknown')})")
+
+            except ImportError:
+                print("Upload requires authentication. Run 'ate login' first.")
+            except Exception as e:
+                print(f"Upload error: {e}")
+
+        return urdf_path
+
+    except KeyboardInterrupt:
+        print("\n\nPipeline interrupted by user")
+        sys.exit(1)
+    except Exception as e:
+        logger.exception("Pipeline failed")
+        print(f"\nPipeline failed: {e}")
+        raise PipelineError(str(e)) from e
+
+
+def resume_pipeline(
+    session_dir: str,
+    from_stage: str = "auto",
+    config: Optional[PipelineConfig] = None,
+) -> Path:
+    """
+    Resume a pipeline from a saved session.
+
+    Args:
+        session_dir: Path to existing session
+        from_stage: Stage to resume from ("capture", "segment", "optimize", "mesh", "synthesize")
+        config: Optional configuration overrides
+
+    Returns:
+        Path to generated URDF
+    """
+    from .scan_session import ScanSession
+
+    print(f"\nResuming pipeline from: {session_dir}")
+
+    # Load session
+    session = ScanSession.load(session_dir)
+    status = session.get_status()
+
+    # Determine resume point
+    if from_stage == "auto":
+        stages = status["stages"]
+        if not stages["capture"]:
+            from_stage = "capture"
+        elif not stages["segment"]:
+            from_stage = "segment"
+        elif not stages["optimize"]:
+            from_stage = "optimize"
+        elif not stages["mesh"]:
+            from_stage = "mesh"
+        else:
+            from_stage = "synthesize"
+
+    print(f"Resuming from: {from_stage}")
+
+    # Create config from session if not provided
+    if config is None:
+        config = PipelineConfig(
+            output_dir=str(session.session_dir),
+            robot_name=session.metadata.robot_name,
+            scale_ref=session.metadata.scale_ref,
+            device=session.metadata.device,
+            density=session.metadata.density_kg_m3,
+        )
+
+    # Run remaining stages
+    segment_data = None
+
+    if from_stage in ["capture"]:
+        session = run_capture_stage(config)
+        from_stage = "segment"
+
+    if from_stage == "segment":
+        segment_data = run_segment_stage(session, config)
+        from_stage = "optimize"
+
+    if from_stage == "optimize":
+        if segment_data is None:
+            # Need to reload segment data
+            print("\nNote: Segment data not in memory, re-running segmentation...")
+            segment_data = run_segment_stage(session, config)
+        run_optimize_stage(session, segment_data, config)
+        from_stage = "mesh"
+
+    if from_stage == "mesh":
+        run_mesh_stage(session, config)
+        from_stage = "synthesize"
+
+    if from_stage == "synthesize":
+        urdf_path = run_synthesize_stage(session, config)
+        return urdf_path
+
+    raise PipelineError(f"Unknown stage: {from_stage}")
+
+
+__all__ = [
+    "PipelineError",
+    "PipelineConfig",
+    "run_capture_stage",
+    "run_segment_stage",
+    "run_optimize_stage",
+    "run_mesh_stage",
+    "run_synthesize_stage",
+    "run_full_pipeline",
+    "resume_pipeline",
+]