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

ate/robot/marker_generator.py

@@ -0,0 +1,611 @@
+#!/usr/bin/env python3
+"""
+ArUco Marker PDF Generator
+===========================
+
+Generates printable ArUco markers for robot calibration.
+Solves the pain point of users needing to figure out marker sizes,
+print settings, and cutting guides.
+
+Usage:
+    ate robot generate-markers --output markers.pdf
+    ate robot generate-markers --count 10 --size 30 --output markers.pdf
+"""
+
+import io
+import math
+from pathlib import Path
+from typing import Optional, List, Tuple
+from dataclasses import dataclass
+
+try:
+    import cv2
+    import numpy as np
+    HAS_OPENCV = True
+except ImportError:
+    HAS_OPENCV = False
+
+try:
+    from reportlab.lib.pagesizes import LETTER, A4
+    from reportlab.lib.units import mm, inch
+    from reportlab.pdfgen import canvas
+    from reportlab.lib.colors import black, gray, white, lightgrey
+    HAS_REPORTLAB = True
+except ImportError:
+    HAS_REPORTLAB = False
+
+
+@dataclass
+class MarkerSpec:
+    """Specification for a single marker."""
+    id: int
+    label: str
+    size_mm: float
+
+
+def check_dependencies():
+    """Check if required dependencies are installed."""
+    missing = []
+    if not HAS_OPENCV:
+        missing.append("opencv-contrib-python")
+    if not HAS_REPORTLAB:
+        missing.append("reportlab")
+
+    if missing:
+        raise ImportError(
+            f"Missing required packages: {', '.join(missing)}. "
+            f"Install with: pip install {' '.join(missing)}"
+        )
+
+
+def generate_aruco_image(
+    marker_id: int,
+    size_pixels: int = 200,
+    dictionary_id: int = cv2.aruco.DICT_4X4_50
+) -> np.ndarray:
+    """Generate an ArUco marker image.
+
+    Args:
+        marker_id: The marker ID (0-49 for DICT_4X4_50)
+        size_pixels: Output image size in pixels
+        dictionary_id: ArUco dictionary to use
+
+    Returns:
+        Numpy array with the marker image (grayscale)
+    """
+    dictionary = cv2.aruco.getPredefinedDictionary(dictionary_id)
+    marker_image = cv2.aruco.generateImageMarker(dictionary, marker_id, size_pixels)
+    return marker_image
+
+
+def generate_marker_pdf(
+    output_path: str,
+    marker_specs: Optional[List[MarkerSpec]] = None,
+    count: int = 12,
+    size_mm: float = 30.0,
+    page_size: Tuple[float, float] = LETTER,
+    include_instructions: bool = True,
+    robot_name: Optional[str] = None,
+) -> str:
+    """Generate a PDF with printable ArUco markers.
+
+    Args:
+        output_path: Path to save the PDF
+        marker_specs: Optional list of marker specifications. If None, generates
+                     generic markers with IDs 0 to count-1.
+        count: Number of markers to generate (ignored if marker_specs provided)
+        size_mm: Marker size in millimeters (ignored if marker_specs provided)
+        page_size: Page size tuple (width, height) in points
+        include_instructions: Include usage instructions on first page
+        robot_name: Optional robot name for labeling
+
+    Returns:
+        Path to the generated PDF
+    """
+    check_dependencies()
+
+    # Default marker specs if not provided
+    if marker_specs is None:
+        marker_specs = [
+            MarkerSpec(id=i, label=f"Joint {i}", size_mm=size_mm)
+            for i in range(count)
+        ]
+
+    # Create PDF
+    c = canvas.Canvas(output_path, pagesize=page_size)
+    page_width, page_height = page_size
+
+    # Margins
+    margin = 0.5 * inch
+    usable_width = page_width - 2 * margin
+    usable_height = page_height - 2 * margin
+
+    # Instructions page
+    if include_instructions:
+        _draw_instructions_page(c, page_width, page_height, robot_name)
+        c.showPage()
+
+    # Calculate grid layout
+    # Add 10mm margin around each marker for cutting
+    marker_with_margin = size_mm + 20  # 10mm on each side
+
+    cols = int(usable_width / (marker_with_margin * mm))
+    rows = int(usable_height / (marker_with_margin * mm))
+
+    if cols < 1:
+        cols = 1
+    if rows < 1:
+        rows = 1
+
+    markers_per_page = cols * rows
+
+    # Center the grid
+    grid_width = cols * marker_with_margin * mm
+    grid_height = rows * marker_with_margin * mm
+    start_x = (page_width - grid_width) / 2 + 10 * mm
+    start_y = page_height - (page_height - grid_height) / 2 - 10 * mm - size_mm * mm
+
+    current_marker = 0
+
+    while current_marker < len(marker_specs):
+        # Draw markers on this page
+        for row in range(rows):
+            if current_marker >= len(marker_specs):
+                break
+            for col in range(cols):
+                if current_marker >= len(marker_specs):
+                    break
+
+                spec = marker_specs[current_marker]
+                x = start_x + col * marker_with_margin * mm
+                y = start_y - row * marker_with_margin * mm
+
+                _draw_marker_cell(c, spec, x, y, spec.size_mm * mm)
+                current_marker += 1
+
+        # Page number
+        c.setFont("Helvetica", 9)
+        c.setFillColor(gray)
+        page_num = (current_marker - 1) // markers_per_page + 1
+        if include_instructions:
+            page_num += 1
+        c.drawCentredString(page_width / 2, 0.3 * inch, f"Page {page_num}")
+
+        if current_marker < len(marker_specs):
+            c.showPage()
+
+    c.save()
+    return output_path
+
+
+def _draw_instructions_page(
+    c: canvas.Canvas,
+    page_width: float,
+    page_height: float,
+    robot_name: Optional[str] = None,
+):
+    """Draw the instructions page."""
+    # Title
+    c.setFont("Helvetica-Bold", 24)
+    c.setFillColor(black)
+    title = "ArUco Calibration Markers"
+    if robot_name:
+        title = f"ArUco Markers for {robot_name}"
+    c.drawCentredString(page_width / 2, page_height - 1 * inch, title)
+
+    # Subtitle
+    c.setFont("Helvetica", 14)
+    c.setFillColor(gray)
+    c.drawCentredString(
+        page_width / 2,
+        page_height - 1.4 * inch,
+        "FoodforThought Robot Calibration System"
+    )
+
+    # Instructions box
+    box_x = 1 * inch
+    box_y = page_height - 6 * inch
+    box_width = page_width - 2 * inch
+    box_height = 4 * inch
+
+    c.setFillColor(lightgrey)
+    c.rect(box_x, box_y, box_width, box_height, fill=True, stroke=False)
+
+    c.setFillColor(black)
+    c.setFont("Helvetica-Bold", 14)
+    c.drawString(box_x + 0.2 * inch, box_y + box_height - 0.4 * inch, "Instructions")
+
+    instructions = [
+        "1. Print this PDF at 100% scale (do not fit to page)",
+        "2. Cut out each marker along the dashed lines",
+        "3. Attach markers to each moving joint of your robot:",
+        "   - Use tape or adhesive putty for temporary mounting",
+        "   - Place marker on rigid part of each joint segment",
+        "   - Ensure markers face the camera and are visible",
+        "4. Run the calibration command:",
+        "   ate robot calibrate --method aruco",
+        "5. Follow the on-screen prompts to wiggle each joint",
+        "6. Upload your calibration to share with the community!",
+    ]
+
+    c.setFont("Helvetica", 11)
+    y = box_y + box_height - 0.8 * inch
+    for line in instructions:
+        c.drawString(box_x + 0.3 * inch, y, line)
+        y -= 0.35 * inch
+
+    # Tips box
+    tips_y = box_y - 0.5 * inch - 2.5 * inch
+    tips_height = 2.5 * inch
+
+    c.setFillColor(white)
+    c.setStrokeColor(gray)
+    c.rect(box_x, tips_y, box_width, tips_height, fill=True, stroke=True)
+
+    c.setFillColor(black)
+    c.setFont("Helvetica-Bold", 12)
+    c.drawString(box_x + 0.2 * inch, tips_y + tips_height - 0.4 * inch, "Tips for Best Results")
+
+    tips = [
+        "- Ensure good, even lighting (avoid shadows and glare)",
+        "- Keep markers flat and perpendicular to camera when possible",
+        "- Use the largest markers that fit on your robot's joints",
+        "- Each marker ID should only appear once on the robot",
+        "- Test detection before running full calibration",
+    ]
+
+    c.setFont("Helvetica", 10)
+    y = tips_y + tips_height - 0.7 * inch
+    for tip in tips:
+        c.drawString(box_x + 0.3 * inch, y, tip)
+        y -= 0.3 * inch
+
+    # Footer
+    c.setFont("Helvetica", 9)
+    c.setFillColor(gray)
+    c.drawCentredString(
+        page_width / 2,
+        0.5 * inch,
+        "Generated by FoodforThought CLI | kindlyrobotics.com"
+    )
+
+
+def _draw_marker_cell(
+    c: canvas.Canvas,
+    spec: MarkerSpec,
+    x: float,
+    y: float,
+    size: float,
+):
+    """Draw a single marker with cutting guides and label."""
+    # Generate marker image
+    marker_img = generate_aruco_image(spec.id, size_pixels=300)
+
+    # Convert to PIL for ReportLab
+    from PIL import Image
+    pil_img = Image.fromarray(marker_img)
+
+    # Save to bytes
+    img_buffer = io.BytesIO()
+    pil_img.save(img_buffer, format='PNG')
+    img_buffer.seek(0)
+
+    # Draw cutting guides (dashed rectangle)
+    padding = 5 * mm
+    cut_x = x - padding
+    cut_y = y - padding
+    cut_size = size + 2 * padding
+
+    c.setStrokeColor(gray)
+    c.setDash(3, 3)
+    c.rect(cut_x, cut_y, cut_size, cut_size, fill=False, stroke=True)
+    c.setDash()  # Reset to solid
+
+    # Corner marks for cutting
+    mark_len = 3 * mm
+    c.setStrokeColor(black)
+    c.setLineWidth(0.5)
+
+    # Top-left corner
+    c.line(cut_x - mark_len, cut_y + cut_size, cut_x, cut_y + cut_size)
+    c.line(cut_x, cut_y + cut_size, cut_x, cut_y + cut_size + mark_len)
+
+    # Top-right corner
+    c.line(cut_x + cut_size, cut_y + cut_size, cut_x + cut_size + mark_len, cut_y + cut_size)
+    c.line(cut_x + cut_size, cut_y + cut_size, cut_x + cut_size, cut_y + cut_size + mark_len)
+
+    # Bottom-left corner
+    c.line(cut_x - mark_len, cut_y, cut_x, cut_y)
+    c.line(cut_x, cut_y, cut_x, cut_y - mark_len)
+
+    # Bottom-right corner
+    c.line(cut_x + cut_size, cut_y, cut_x + cut_size + mark_len, cut_y)
+    c.line(cut_x + cut_size, cut_y, cut_x + cut_size, cut_y - mark_len)
+
+    # Draw the marker image
+    from reportlab.lib.utils import ImageReader
+    img_reader = ImageReader(img_buffer)
+    c.drawImage(img_reader, x, y, width=size, height=size)
+
+    # Label below marker
+    c.setFont("Helvetica", 8)
+    c.setFillColor(black)
+    label = f"ID: {spec.id}"
+    if spec.label:
+        label = f"{spec.label} (ID: {spec.id})"
+    c.drawCentredString(x + size / 2, y - 4 * mm, label)
+
+    # Size indicator
+    c.setFont("Helvetica", 6)
+    c.setFillColor(gray)
+    c.drawCentredString(x + size / 2, y - 7 * mm, f"{spec.size_mm:.0f}mm")
+
+
+def generate_robot_markers(
+    robot_joints: List[dict],
+    output_path: str,
+    robot_name: str,
+    base_size_mm: float = 25.0,
+) -> str:
+    """Generate markers specifically for a robot's joints.
+
+    Args:
+        robot_joints: List of joint dicts with 'name', 'type', optional 'segment'
+        output_path: Path to save the PDF
+        robot_name: Name of the robot
+        base_size_mm: Base marker size in mm
+
+    Returns:
+        Path to the generated PDF
+    """
+    # Create marker specs from joint definitions
+    specs = []
+    for i, joint in enumerate(robot_joints):
+        name = joint.get('name', f'Joint {i}')
+        segment = joint.get('segment', '')
+
+        label = name
+        if segment:
+            label = f"{segment}: {name}"
+
+        specs.append(MarkerSpec(
+            id=i,
+            label=label,
+            size_mm=base_size_mm,
+        ))
+
+    return generate_marker_pdf(
+        output_path=output_path,
+        marker_specs=specs,
+        robot_name=robot_name,
+        include_instructions=True,
+    )
+
+
+def generate_quick_markers(
+    output_path: str = "aruco_markers.pdf",
+    count: int = 12,
+    size_mm: float = 30.0,
+) -> str:
+    """Quick generation of generic calibration markers.
+
+    Args:
+        output_path: Path to save the PDF
+        count: Number of markers to generate
+        size_mm: Size of each marker in millimeters
+
+    Returns:
+        Path to the generated PDF
+    """
+    return generate_marker_pdf(
+        output_path=output_path,
+        count=count,
+        size_mm=size_mm,
+        include_instructions=True,
+    )
+
+
+# Robot-specific marker presets
+# These encode expert knowledge about marker placement and sizes
+ROBOT_MARKER_PRESETS = {
+    "mechdog-mini": {
+        "robot_name": "Hiwonder MechDog Mini",
+        "markers": [
+            # Body - large, visible from distance
+            MarkerSpec(id=0, label="Body Center", size_mm=35),
+            # Front Right Leg
+            MarkerSpec(id=1, label="FR Hip", size_mm=25),
+            MarkerSpec(id=2, label="FR Thigh", size_mm=25),
+            MarkerSpec(id=3, label="FR Shin", size_mm=20),
+            # Front Left Leg
+            MarkerSpec(id=4, label="FL Hip", size_mm=25),
+            MarkerSpec(id=5, label="FL Thigh", size_mm=25),
+            MarkerSpec(id=6, label="FL Shin", size_mm=20),
+            # Back Right Leg
+            MarkerSpec(id=7, label="BR Hip", size_mm=25),
+            MarkerSpec(id=8, label="BR Thigh", size_mm=25),
+            MarkerSpec(id=9, label="BR Shin", size_mm=20),
+            # Back Left Leg
+            MarkerSpec(id=10, label="BL Hip", size_mm=25),
+            MarkerSpec(id=11, label="BL Thigh", size_mm=25),
+            MarkerSpec(id=12, label="BL Shin", size_mm=20),
+            # Arm (servo 11, 12, 13 from our earlier discovery)
+            MarkerSpec(id=13, label="Arm Shoulder", size_mm=25),
+            MarkerSpec(id=14, label="Arm Elbow", size_mm=20),
+            MarkerSpec(id=15, label="Gripper", size_mm=15),
+        ],
+    },
+    "unitree-go1": {
+        "robot_name": "Unitree Go1",
+        "markers": [
+            MarkerSpec(id=0, label="Body Center", size_mm=60),
+            # Larger robot = larger markers
+            MarkerSpec(id=1, label="FR Hip", size_mm=40),
+            MarkerSpec(id=2, label="FR Thigh", size_mm=40),
+            MarkerSpec(id=3, label="FR Calf", size_mm=35),
+            MarkerSpec(id=4, label="FL Hip", size_mm=40),
+            MarkerSpec(id=5, label="FL Thigh", size_mm=40),
+            MarkerSpec(id=6, label="FL Calf", size_mm=35),
+            MarkerSpec(id=7, label="RR Hip", size_mm=40),
+            MarkerSpec(id=8, label="RR Thigh", size_mm=40),
+            MarkerSpec(id=9, label="RR Calf", size_mm=35),
+            MarkerSpec(id=10, label="RL Hip", size_mm=40),
+            MarkerSpec(id=11, label="RL Thigh", size_mm=40),
+            MarkerSpec(id=12, label="RL Calf", size_mm=35),
+        ],
+    },
+    "xarm-6": {
+        "robot_name": "xArm 6-DOF",
+        "markers": [
+            MarkerSpec(id=0, label="Base", size_mm=50),
+            MarkerSpec(id=1, label="Shoulder", size_mm=45),
+            MarkerSpec(id=2, label="Upper Arm", size_mm=40),
+            MarkerSpec(id=3, label="Forearm", size_mm=40),
+            MarkerSpec(id=4, label="Wrist 1", size_mm=35),
+            MarkerSpec(id=5, label="Wrist 2", size_mm=35),
+            MarkerSpec(id=6, label="End Effector", size_mm=30),
+        ],
+    },
+    "generic-quadruped": {
+        "robot_name": "Generic Quadruped",
+        "markers": [
+            MarkerSpec(id=0, label="Body", size_mm=40),
+            MarkerSpec(id=1, label="FR Segment 1", size_mm=30),
+            MarkerSpec(id=2, label="FR Segment 2", size_mm=25),
+            MarkerSpec(id=3, label="FR Segment 3", size_mm=20),
+            MarkerSpec(id=4, label="FL Segment 1", size_mm=30),
+            MarkerSpec(id=5, label="FL Segment 2", size_mm=25),
+            MarkerSpec(id=6, label="FL Segment 3", size_mm=20),
+            MarkerSpec(id=7, label="BR Segment 1", size_mm=30),
+            MarkerSpec(id=8, label="BR Segment 2", size_mm=25),
+            MarkerSpec(id=9, label="BR Segment 3", size_mm=20),
+            MarkerSpec(id=10, label="BL Segment 1", size_mm=30),
+            MarkerSpec(id=11, label="BL Segment 2", size_mm=25),
+            MarkerSpec(id=12, label="BL Segment 3", size_mm=20),
+        ],
+    },
+    "generic-arm": {
+        "robot_name": "Generic Robot Arm",
+        "markers": [
+            MarkerSpec(id=0, label="Base", size_mm=45),
+            MarkerSpec(id=1, label="Joint 1", size_mm=40),
+            MarkerSpec(id=2, label="Joint 2", size_mm=35),
+            MarkerSpec(id=3, label="Joint 3", size_mm=35),
+            MarkerSpec(id=4, label="Joint 4", size_mm=30),
+            MarkerSpec(id=5, label="Joint 5", size_mm=30),
+            MarkerSpec(id=6, label="End Effector", size_mm=25),
+        ],
+    },
+}
+
+
+def get_preset_markers(preset_name: str) -> Tuple[List[MarkerSpec], str]:
+    """Get marker specs for a known robot type.
+
+    Args:
+        preset_name: Robot preset name (e.g., "mechdog-mini", "unitree-go1")
+
+    Returns:
+        Tuple of (marker_specs, robot_name)
+
+    Raises:
+        KeyError if preset not found
+    """
+    if preset_name not in ROBOT_MARKER_PRESETS:
+        available = ", ".join(ROBOT_MARKER_PRESETS.keys())
+        raise KeyError(f"Unknown preset '{preset_name}'. Available: {available}")
+
+    preset = ROBOT_MARKER_PRESETS[preset_name]
+    return preset["markers"], preset["robot_name"]
+
+
+def list_presets() -> List[dict]:
+    """List available robot presets."""
+    return [
+        {
+            "name": name,
+            "robot_name": preset["robot_name"],
+            "marker_count": len(preset["markers"]),
+        }
+        for name, preset in ROBOT_MARKER_PRESETS.items()
+    ]
+
+
+# CLI integration
+def add_marker_command(subparsers):
+    """Add the generate-markers command to the CLI."""
+    parser = subparsers.add_parser(
+        "generate-markers",
+        help="Generate printable ArUco markers for robot calibration",
+    )
+    parser.add_argument(
+        "--output", "-o",
+        default="aruco_markers.pdf",
+        help="Output PDF path (default: aruco_markers.pdf)",
+    )
+    parser.add_argument(
+        "--count", "-n",
+        type=int,
+        default=12,
+        help="Number of markers to generate (default: 12)",
+    )
+    parser.add_argument(
+        "--size",
+        type=float,
+        default=30.0,
+        help="Marker size in millimeters (default: 30)",
+    )
+    parser.add_argument(
+        "--robot-name",
+        help="Optional robot name for labeling",
+    )
+    parser.add_argument(
+        "--page-size",
+        choices=["letter", "a4"],
+        default="letter",
+        help="Page size (default: letter)",
+    )
+
+    return parser
+
+
+def run_marker_command(args):
+    """Run the generate-markers command."""
+    from reportlab.lib.pagesizes import LETTER, A4
+
+    page_size = LETTER if args.page_size == "letter" else A4
+
+    print(f"Generating {args.count} ArUco markers...")
+    print(f"  Size: {args.size}mm")
+    print(f"  Output: {args.output}")
+
+    output = generate_marker_pdf(
+        output_path=args.output,
+        count=args.count,
+        size_mm=args.size,
+        page_size=page_size,
+        robot_name=args.robot_name,
+        include_instructions=True,
+    )
+
+    print(f"\n[OK] Markers saved to: {output}")
+    print("\nNext steps:")
+    print("  1. Print the PDF at 100% scale")
+    print("  2. Cut out markers along the dashed lines")
+    print("  3. Attach to your robot's joints")
+    print("  4. Run: ate robot calibrate --method aruco")
+
+
+if __name__ == "__main__":
+    # Quick test
+    import sys
+
+    output = "/tmp/test_markers.pdf"
+    print(f"Generating test markers to {output}...")
+
+    try:
+        generate_quick_markers(output, count=6, size_mm=40)
+        print(f"Done! Open {output} to view.")
+    except ImportError as e:
+        print(f"Error: {e}")
+        sys.exit(1)