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

ate/urdf/commands.py

@@ -0,0 +1,708 @@
+"""
+CLI command definitions for URDF scan pipeline.
+
+Provides the following commands:
+- ate urdf scan capture    - Capture video and annotate links
+- ate urdf scan segment    - Generate segmented point clouds
+- ate urdf scan optimize   - Estimate joint parameters
+- ate urdf scan mesh       - Generate visual/collision meshes
+- ate urdf scan synthesize - Generate final URDF
+- ate urdf scan            - Run full pipeline
+- ate urdf scan status     - Show session status
+"""
+
+import sys
+import argparse
+from pathlib import Path
+from typing import Optional
+
+
+def _has_mps() -> bool:
+    """Check if Apple Metal Performance Shaders (MPS) is available."""
+    try:
+        import torch
+        return torch.backends.mps.is_available()
+    except (ImportError, AttributeError):
+        return False
+
+def register_parser(subparsers):
+    """Register URDF scan commands with argparse."""
+    urdf_parser = subparsers.add_parser(
+        "urdf",
+        help="URDF generation and management",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+
+    urdf_subparsers = urdf_parser.add_subparsers(
+        dest="urdf_action",
+        help="URDF action"
+    )
+
+    # urdf scan - parent for scan subcommands
+    scan_parser = urdf_subparsers.add_parser(
+        "scan",
+        help="Markerless URDF generation from webcam video",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        description="""
+Markerless URDF Generation Pipeline
+
+Generates simulation-ready URDF files from webcam video without
+physical markers, using foundation vision models (SAM 2, Depth Anything V2).
+
+QUICK START:
+    ate urdf scan --output ./my_robot/ --scale-ref "gripper:85mm"
+
+STEP-BY-STEP:
+    ate urdf scan capture --output ./my_robot/
+    ate urdf scan segment ./my_robot/ --scale-ref "gripper:85mm"
+    ate urdf scan optimize ./my_robot/
+    ate urdf scan mesh ./my_robot/
+    ate urdf scan synthesize ./my_robot/ --name my_robot
+"""
+    )
+
+    scan_subparsers = scan_parser.add_subparsers(
+        dest="scan_action",
+        help="Scan action"
+    )
+
+    # Common arguments
+    def add_common_args(parser):
+        parser.add_argument(
+            "--device", choices=["cpu", "cuda"], default="cpu",
+            help="Compute device (default: cpu)"
+        )
+        parser.add_argument(
+            "--dry-run", action="store_true",
+            help="Preview operations without executing"
+        )
+
+    # urdf scan (full pipeline)
+    scan_parser.add_argument(
+        "-o", "--output", default="./urdf_output",
+        help="Output directory for scan session (default: ./urdf_output)"
+    )
+    scan_parser.add_argument(
+        "-n", "--name",
+        help="Robot name (default: directory name)"
+    )
+    scan_parser.add_argument(
+        "-s", "--scale-ref",
+        help="Scale reference (e.g., 'gripper:85mm', 'base_width:150mm')"
+    )
+    scan_parser.add_argument(
+        "-v", "--video",
+        help="Path to existing video (skip capture)"
+    )
+    scan_parser.add_argument(
+        "--camera", type=int, default=0,
+        help="Camera device ID (default: 0)"
+    )
+    scan_parser.add_argument(
+        "--density", type=float, default=1200.0,
+        help="Material density in kg/m^3 (default: 1200, plastic)"
+    )
+    scan_parser.add_argument(
+        "--upload", action="store_true",
+        help="Upload to FoodforThought platform after generation"
+    )
+    scan_parser.add_argument(
+        "--cloud", action="store_true",
+        help="Process scan on cloud GPU (uses credits, no local GPU required)"
+    )
+    scan_parser.add_argument(
+        "--local", action="store_true",
+        help="Force local processing (requires GPU with 8GB+ VRAM or Apple Silicon)"
+    )
+    add_common_args(scan_parser)
+
+    # urdf scan capture
+    capture_parser = scan_subparsers.add_parser(
+        "capture",
+        help="Capture video and annotate robot links"
+    )
+    capture_parser.add_argument(
+        "-o", "--output", required=True,
+        help="Output directory for scan session"
+    )
+    capture_parser.add_argument(
+        "-v", "--video",
+        help="Path to existing video (skip webcam capture)"
+    )
+    capture_parser.add_argument(
+        "--camera", type=int, default=0,
+        help="Camera device ID (default: 0)"
+    )
+    capture_parser.add_argument(
+        "-n", "--name",
+        help="Robot name"
+    )
+    capture_parser.add_argument(
+        "-s", "--scale-ref",
+        help="Scale reference"
+    )
+    add_common_args(capture_parser)
+
+    # urdf scan segment
+    segment_parser = scan_subparsers.add_parser(
+        "segment",
+        help="Generate segmented point clouds from video"
+    )
+    segment_parser.add_argument(
+        "session_dir",
+        help="Path to scan session directory"
+    )
+    segment_parser.add_argument(
+        "-s", "--scale-ref",
+        help="Scale reference (e.g., 'gripper:85mm')"
+    )
+    segment_parser.add_argument(
+        "--frame-skip", type=int, default=1,
+        help="Process every Nth frame (default: 1)"
+    )
+    segment_parser.add_argument(
+        "--voxel-size", type=float, default=0.005,
+        help="Voxel size for point cloud downsampling (meters, default: 0.005)"
+    )
+    add_common_args(segment_parser)
+
+    # urdf scan optimize
+    optimize_parser = scan_subparsers.add_parser(
+        "optimize",
+        help="Estimate kinematic joint parameters"
+    )
+    optimize_parser.add_argument(
+        "session_dir",
+        help="Path to scan session directory"
+    )
+    optimize_parser.add_argument(
+        "--joint-types",
+        help="Override joint types (e.g., 'shoulder:revolute,gripper:prismatic')"
+    )
+    add_common_args(optimize_parser)
+
+    # urdf scan mesh
+    mesh_parser = scan_subparsers.add_parser(
+        "mesh",
+        help="Generate visual and collision meshes"
+    )
+    mesh_parser.add_argument(
+        "session_dir",
+        help="Path to scan session directory"
+    )
+    mesh_parser.add_argument(
+        "--visual-only", action="store_true",
+        help="Skip collision mesh generation"
+    )
+    mesh_parser.add_argument(
+        "--collision-hulls", type=int, default=8,
+        help="Max convex hulls per link (default: 8)"
+    )
+    mesh_parser.add_argument(
+        "--simplify-to", type=int, default=5000,
+        help="Target face count for mesh simplification (default: 5000)"
+    )
+    add_common_args(mesh_parser)
+
+    # urdf scan synthesize
+    synthesize_parser = scan_subparsers.add_parser(
+        "synthesize",
+        help="Generate final URDF file"
+    )
+    synthesize_parser.add_argument(
+        "session_dir",
+        help="Path to scan session directory"
+    )
+    synthesize_parser.add_argument(
+        "-n", "--name",
+        help="Robot name in URDF"
+    )
+    synthesize_parser.add_argument(
+        "--density", type=float, default=1200.0,
+        help="Material density for inertial estimation (kg/m^3)"
+    )
+    synthesize_parser.add_argument(
+        "--upload", action="store_true",
+        help="Upload to FoodforThought platform"
+    )
+    add_common_args(synthesize_parser)
+
+    # urdf scan status
+    status_parser = scan_subparsers.add_parser(
+        "status",
+        help="Show scan session status"
+    )
+    status_parser.add_argument(
+        "session_dir",
+        help="Path to scan session directory"
+    )
+    status_parser.add_argument(
+        "--json", action="store_true", dest="json_output",
+        help="Output as JSON"
+    )
+
+    # urdf scan validate
+    validate_parser = scan_subparsers.add_parser(
+        "validate",
+        help="Validate generated URDF"
+    )
+    validate_parser.add_argument(
+        "urdf_path",
+        help="Path to URDF file"
+    )
+    validate_parser.add_argument(
+        "--no-mesh-check", action="store_true",
+        help="Skip mesh file existence checks"
+    )
+
+    # urdf scan jobs - list cloud scan jobs
+    jobs_parser = scan_subparsers.add_parser(
+        "jobs",
+        help="List cloud scan jobs"
+    )
+    jobs_parser.add_argument(
+        "--json", action="store_true", dest="json_output",
+        help="Output as JSON"
+    )
+    jobs_parser.add_argument(
+        "--limit", type=int, default=10,
+        help="Maximum jobs to display (default: 10)"
+    )
+
+    # urdf scan job - show details of a specific cloud job
+    job_parser = scan_subparsers.add_parser(
+        "job",
+        help="Show details of a cloud scan job"
+    )
+    job_parser.add_argument(
+        "job_id",
+        help="Job ID to check"
+    )
+    job_parser.add_argument(
+        "--json", action="store_true", dest="json_output",
+        help="Output as JSON"
+    )
+
+    # urdf scan download - download artifact from completed cloud job
+    download_parser = scan_subparsers.add_parser(
+        "download",
+        help="Download artifact from completed cloud scan job"
+    )
+    download_parser.add_argument(
+        "job_id",
+        help="Job ID to download artifact from"
+    )
+    download_parser.add_argument(
+        "--output", "-o", default=".",
+        help="Output directory (default: current directory)"
+    )
+
+
+def handle(client, args):
+    """Handle URDF commands."""
+    if args.urdf_action == "scan":
+        handle_scan(args)
+    else:
+        print(f"Unknown URDF action: {args.urdf_action}")
+        print("Available actions: scan")
+        sys.exit(1)
+
+
+def handle_scan(args):
+    """Handle urdf scan commands."""
+    from .pipeline import PipelineConfig, run_full_pipeline, resume_pipeline
+    from .scan_session import ScanSession
+
+    # If no scan action, run full pipeline
+    if args.scan_action is None:
+        # Check for cloud mode
+        use_cloud = getattr(args, 'cloud', False)
+        use_local = getattr(args, 'local', False)
+
+        if use_cloud and use_local:
+            print("Error: Cannot specify both --cloud and --local")
+            sys.exit(1)
+
+        if use_cloud:
+            # Run cloud pipeline
+            from .cloud import run_cloud_pipeline
+            try:
+                run_cloud_pipeline(
+                    video_path=getattr(args, 'video', None),
+                    output_dir=args.output,
+                    robot_name=getattr(args, 'name', None),
+                    scale_ref=getattr(args, 'scale_ref', None),
+                    camera_id=getattr(args, 'camera', 0),
+                )
+            except RuntimeError as e:
+                print(f"\nCloud processing failed: {e}")
+                sys.exit(1)
+            return
+
+        # Local pipeline (default)
+        # Check hardware capabilities if not forced
+        if not use_local:
+            # Recommend cloud if local hardware is insufficient
+            try:
+                import torch
+                if not torch.cuda.is_available() and not _has_mps():
+                    print("\nNote: No GPU detected. Consider using --cloud for faster processing.")
+                    print("      Cloud processing costs 5 credits per scan.")
+                    print("      Run 'ate urdf scan --cloud' to use cloud GPU.\n")
+            except ImportError:
+                pass
+
+        # Full local pipeline
+        config = PipelineConfig(
+            output_dir=args.output,
+            robot_name=getattr(args, 'name', None),
+            scale_ref=getattr(args, 'scale_ref', None),
+            video_path=getattr(args, 'video', None),
+            device=getattr(args, 'device', 'cpu'),
+            camera_id=getattr(args, 'camera', 0),
+            density=getattr(args, 'density', 1200.0),
+            dry_run=getattr(args, 'dry_run', False),
+            upload=getattr(args, 'upload', False),
+        )
+        run_full_pipeline(config)
+
+    elif args.scan_action == "capture":
+        from .capture import run_capture
+        session = run_capture(
+            output_dir=args.output,
+            video_path=getattr(args, 'video', None),
+            camera_id=getattr(args, 'camera', 0),
+            robot_name=getattr(args, 'name', None),
+            scale_ref=getattr(args, 'scale_ref', None),
+        )
+        print(f"\nSession created: {session.session_dir}")
+
+    elif args.scan_action == "segment":
+        from .segmentation import run_segmentation
+        from .depth import run_depth_estimation
+        from .lifting import run_lifting
+
+        session = ScanSession.load(args.session_dir)
+
+        # Update scale ref if provided
+        if args.scale_ref:
+            session.metadata.scale_ref = args.scale_ref
+            session.save_metadata()
+
+        print("Running segmentation...")
+        masks = run_segmentation(session)
+
+        print("Running depth estimation...")
+        depth_maps, _ = run_depth_estimation(
+            session, masks,
+            frame_skip=args.frame_skip,
+        )
+
+        print("Generating point clouds...")
+        run_lifting(
+            session, masks, depth_maps,
+            frame_skip=args.frame_skip,
+            voxel_size=args.voxel_size,
+        )
+
+        print(f"\nSegmentation complete: {session.session_dir}")
+
+    elif args.scan_action == "optimize":
+        from .kinematics import run_kinematic_optimization
+        from .segmentation import run_segmentation
+        from .depth import run_depth_estimation
+
+        session = ScanSession.load(args.session_dir)
+
+        # Parse joint type hints
+        joint_hints = None
+        if args.joint_types:
+            joint_hints = {}
+            for pair in args.joint_types.split(","):
+                link, jtype = pair.split(":")
+                joint_hints[link.strip()] = jtype.strip()
+
+        # Need masks and depth_maps - re-run if needed
+        print("Running segmentation...")
+        masks = run_segmentation(session)
+
+        print("Running depth estimation...")
+        depth_maps, _ = run_depth_estimation(session, masks)
+
+        print("Optimizing kinematics...")
+        joints = run_kinematic_optimization(
+            session, masks, depth_maps,
+            joint_type_hints=joint_hints,
+        )
+
+        print(f"\nOptimization complete: {len(joints)} joints")
+
+    elif args.scan_action == "mesh":
+        from .meshing import generate_all_visual_meshes
+        from .collision import generate_all_collision_meshes
+
+        session = ScanSession.load(args.session_dir)
+
+        print("Generating visual meshes...")
+        visual_paths = generate_all_visual_meshes(
+            session,
+            simplify_to=args.simplify_to,
+        )
+
+        if not args.visual_only:
+            print("Generating collision meshes...")
+            generate_all_collision_meshes(
+                session,
+                max_hulls=args.collision_hulls,
+            )
+
+        print(f"\nMesh generation complete: {len(visual_paths)} links")
+
+    elif args.scan_action == "synthesize":
+        from .synthesis import run_synthesis
+        from .validation import run_validation
+
+        session = ScanSession.load(args.session_dir)
+
+        print("Synthesizing URDF...")
+        urdf_path = run_synthesis(
+            session,
+            density=args.density,
+            robot_name=getattr(args, 'name', None),
+        )
+
+        print("\nValidating URDF...")
+        result = run_validation(session)
+        print(result.summary())
+
+        if args.upload:
+            print("\nUploading to FoodforThought...")
+            try:
+                from ..robot.skill_upload import SkillLibraryUploader, APIError
+
+                uploader = SkillLibraryUploader()
+                robot_name = session.metadata.robot_name or "unknown_robot"
+                
+                # Get or create project for this robot
+                project = uploader.get_or_create_project(
+                    name=f"{robot_name}_urdf",
+                    description=f"URDF generated from markerless scan for {robot_name}",
+                )
+                project_id = project["id"]
+
+                # Upload URDF as artifact
+                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,
+                    },
+                )
+
+                artifact_id = response.get("id", "unknown")
+                print(f"Uploaded URDF: {robot_name}.urdf (artifact {artifact_id})")
+                print(f"Project: {project.get('name', project_id)}")
+
+            except ImportError:
+                print("Upload requires authentication. Run 'ate login' first.")
+            except APIError as e:
+                print(f"Upload failed: {e}")
+            except Exception as e:
+                print(f"Upload error: {e}")
+
+        print(f"\nGenerated: {urdf_path}")
+
+    elif args.scan_action == "status":
+        import json
+        session = ScanSession.load(args.session_dir)
+        status = session.get_status()
+
+        if args.json_output:
+            print(json.dumps(status, indent=2))
+        else:
+            print(f"\nSession: {status['session_dir']}")
+            print(f"Robot: {status['robot_name']}")
+            print(f"Scale: {status['scale_ref'] or 'not set'}")
+            print(f"\nStages:")
+            for stage, complete in status['stages'].items():
+                check = "[x]" if complete else "[ ]"
+                print(f"  {check} {stage}")
+            print(f"\nData:")
+            print(f"  Video: {'yes' if status['data']['has_video'] else 'no'}")
+            print(f"  Links: {status['data']['link_count']}")
+            print(f"  Joints: {status['data']['joint_count']}")
+            print(f"  Point clouds: {'yes' if status['data']['has_clouds'] else 'no'}")
+            print(f"  Meshes: {'yes' if status['data']['has_meshes'] else 'no'}")
+            print(f"  URDF: {'yes' if status['data']['has_urdf'] else 'no'}")
+
+    elif args.scan_action == "validate":
+        from .validation import validate_urdf
+
+        urdf_path = Path(args.urdf_path)
+        result = validate_urdf(urdf_path, check_meshes=not args.no_mesh_check)
+        print(result.summary())
+        sys.exit(0 if result.valid else 1)
+
+    elif args.scan_action == "jobs":
+        import json
+        from datetime import datetime
+        from .cloud import CloudScanClient
+
+        client = CloudScanClient()
+        result = client.list_jobs()
+
+        if "error" in result:
+            print(f"Error: {result['error']}")
+            sys.exit(1)
+
+        jobs = result.get("jobs", [])
+        credits = result.get("credits", {})
+
+        if args.json_output:
+            print(json.dumps(result, indent=2))
+        else:
+            print(f"\nCloud Scan Jobs (Credits: {credits.get('balance', 0)})")
+            print("=" * 70)
+
+            if not jobs:
+                print("No jobs found. Run 'ate urdf scan --cloud' to create one.")
+            else:
+                # Show most recent first, up to limit
+                jobs = sorted(jobs, key=lambda j: j.get("createdAt", ""), reverse=True)
+                jobs = jobs[:args.limit]
+
+                for job in jobs:
+                    job_id = job.get("id", "unknown")[:25]
+                    status = job.get("status", "unknown")
+                    robot = job.get("robotName") or "-"
+                    created = job.get("createdAt", "")
+                    
+                    # Format relative time
+                    time_str = ""
+                    if created:
+                        try:
+                            dt = datetime.fromisoformat(created.replace("Z", "+00:00"))
+                            diff = datetime.now(dt.tzinfo) - dt
+                            if diff.days > 0:
+                                time_str = f"{diff.days}d ago"
+                            elif diff.seconds > 3600:
+                                time_str = f"{diff.seconds // 3600}h ago"
+                            else:
+                                time_str = f"{diff.seconds // 60}m ago"
+                        except Exception:
+                            time_str = created[:10]
+
+                    # Status emoji
+                    status_icon = {
+                        "completed": "✓",
+                        "processing": "⏳",
+                        "uploading": "↑",
+                        "pending": "○",
+                        "failed": "✗"
+                    }.get(status, "?")
+
+                    print(f"{status_icon} {job_id}  {status:12}  {robot:12}  {time_str}")
+
+            print()
+
+    elif args.scan_action == "job":
+        import json
+        from .cloud import CloudScanClient
+
+        client = CloudScanClient()
+        job = client.get_job(args.job_id)
+
+        if "error" in job:
+            print(f"Error: {job['error']}")
+            sys.exit(1)
+
+        if args.json_output:
+            print(json.dumps(job, indent=2))
+        else:
+            status = job.get("status", "unknown")
+            status_icon = {
+                "completed": "✓",
+                "processing": "⏳",
+                "uploading": "↑",
+                "pending": "○",
+                "failed": "✗"
+            }.get(status, "?")
+
+            print(f"\nJob: {job.get('id', 'unknown')}")
+            print(f"Status: {status_icon} {status}")
+            print(f"Robot: {job.get('robotName') or '-'}")
+            print(f"Scale: {job.get('scaleRef') or '-'}")
+            print(f"Created: {job.get('createdAt', '-')}")
+
+            if status == "completed":
+                print(f"Links: {job.get('linkCount', '-')}")
+                print(f"Joints: {job.get('jointCount', '-')}")
+                print(f"Processing: {job.get('processingTimeSeconds', '-')}s")
+                if job.get("artifactUrl"):
+                    print(f"Artifact: {job['artifactUrl'][:60]}...")
+                print(f"\nDownload: ate urdf scan download {job.get('id')}")
+
+            elif status == "failed":
+                print(f"Error: {job.get('error', 'Unknown error')}")
+
+            print()
+
+    elif args.scan_action == "download":
+        from pathlib import Path
+        from .cloud import CloudScanClient
+
+        client = CloudScanClient()
+        job = client.get_job(args.job_id)
+
+        if "error" in job:
+            print(f"Error: {job['error']}")
+            sys.exit(1)
+
+        if job.get("status") != "completed":
+            print(f"Error: Job is '{job.get('status')}', not 'completed'")
+            print("Wait for processing to complete before downloading.")
+            sys.exit(1)
+
+        artifact_url = job.get("artifactUrl")
+        if not artifact_url:
+            print("Error: No artifact URL found for this job")
+            sys.exit(1)
+
+        output_dir = Path(args.output)
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        print(f"Downloading artifact from job {args.job_id}...")
+
+        def download_progress(downloaded, total):
+            pct = int(100 * downloaded / total) if total > 0 else 0
+            print(f"\r  {pct}% ({downloaded // 1024} KB)", end="", flush=True)
+
+        try:
+            urdf_path = client.download_artifact(artifact_url, output_dir, download_progress)
+            print(f"\n\nDownloaded to: {output_dir}/")
+            print(f"URDF: {urdf_path}")
+        except RuntimeError as e:
+            print(f"\nError: {e}")
+            sys.exit(1)
+
+    else:
+        print(f"Unknown scan action: {args.scan_action}")
+        print("Available actions: capture, segment, optimize, mesh, synthesize, status, validate, jobs, job, download")
+        sys.exit(1)
+
+
+__all__ = [
+    "register_parser",
+    "handle",
+]