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

ate/urdf/capture.py

@@ -0,0 +1,582 @@
+"""
+Video capture and link annotation for URDF generation.
+
+This module provides the interactive capture UI for Phase 1 of the
+markerless URDF pipeline:
+
+1. Open webcam feed with overlay guidance
+2. Capture video of robot moving through range of motion
+3. Pause on first frame for link annotation
+4. Save video, link annotations, and metadata
+
+Usage:
+    ate urdf scan capture --output ./my_robot_scan/
+    ate urdf scan capture --video ./existing_video.mp4 --output ./my_robot_scan/
+"""
+
+import sys
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List, Optional, Tuple, Callable
+import logging
+
+logger = logging.getLogger(__name__)
+
+try:
+    import cv2
+    import numpy as np
+    CV2_AVAILABLE = True
+except ImportError:
+    CV2_AVAILABLE = False
+    cv2 = None
+    np = None
+
+
+class CaptureError(Exception):
+    """Error during video capture."""
+    pass
+
+
+@dataclass
+class CaptureConfig:
+    """Configuration for video capture."""
+    camera_id: int = 0
+    width: int = 1280
+    height: int = 720
+    fps: float = 30.0
+    codec: str = "mp4v"
+    flip_horizontal: bool = False
+
+
+class LinkAnnotator:
+    """
+    Interactive link annotation on a video frame.
+
+    Allows user to click on robot links to define segmentation points.
+    """
+
+    def __init__(self, frame: "np.ndarray"):
+        """
+        Initialize annotator with a frame.
+
+        Args:
+            frame: BGR image from OpenCV
+        """
+        if not CV2_AVAILABLE:
+            raise CaptureError("OpenCV not installed. Run: pip install opencv-python")
+
+        self.frame = frame.copy()
+        self.display_frame = frame.copy()
+        self.links: List[Tuple[str, List[float], bool]] = []  # (name, [x,y], is_fixed)
+        self.current_point: Optional[Tuple[int, int]] = None
+        self.window_name = "URDF Scan - Link Annotation"
+        self._running = True
+        self._last_click_time = 0.0  # Debounce for double-click handling
+
+    def _draw_overlay(self) -> "np.ndarray":
+        """Draw annotation overlay on frame."""
+        display = self.display_frame.copy()
+        h, w = display.shape[:2]
+
+        # Draw existing link points
+        for i, (name, point, is_fixed) in enumerate(self.links):
+            x, y = int(point[0]), int(point[1])
+            color = (0, 255, 0) if is_fixed else (255, 165, 0)  # Green for fixed, orange for movable
+            cv2.circle(display, (x, y), 8, color, -1)
+            cv2.circle(display, (x, y), 10, (255, 255, 255), 2)
+
+            # Label
+            label = f"{i+1}. {name}"
+            if is_fixed:
+                label += " (base)"
+            cv2.putText(
+                display, label, (x + 15, y + 5),
+                cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2
+            )
+            cv2.putText(
+                display, label, (x + 15, y + 5),
+                cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 1
+            )
+
+        # Draw current point if hovering
+        if self.current_point:
+            x, y = self.current_point
+            cv2.circle(display, (x, y), 8, (100, 100, 255), 2)
+
+        # Instructions panel
+        panel_h = 150
+        overlay = display.copy()
+        cv2.rectangle(overlay, (10, 10), (w - 10, panel_h), (0, 0, 0), -1)
+        cv2.addWeighted(overlay, 0.7, display, 0.3, 0, display)
+
+        # Instructions text
+        instructions = [
+            "LINK ANNOTATION - Click on each rigid robot part",
+            "",
+            f"Links annotated: {len(self.links)}",
+            "",
+            "Controls:",
+            "  [Click] Add link point    [Enter/Space] Finish",
+            "  [U] Undo last point       [Esc] Cancel",
+            "  [B] Mark last as base (fixed link)",
+        ]
+
+        y_offset = 30
+        for line in instructions:
+            cv2.putText(
+                display, line, (20, y_offset),
+                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1
+            )
+            y_offset += 18
+
+        return display
+
+    def _mouse_callback(self, event, x, y, flags, param):
+        """Handle mouse events."""
+        if event == cv2.EVENT_MOUSEMOVE:
+            self.current_point = (x, y)
+        elif event == cv2.EVENT_LBUTTONDOWN:
+            # Debounce: ignore clicks within 300ms of previous click
+            current_time = time.time()
+            if current_time - self._last_click_time < 0.3:
+                return  # Ignore this click (too fast, likely double-click artifact)
+            self._last_click_time = current_time
+            self._add_link_at(x, y)
+
+    def _add_link_at(self, x: int, y: int) -> None:
+        """Add a link annotation at the given coordinates."""
+        # Prompt for link name
+        link_num = len(self.links) + 1
+        default_names = ["base", "shoulder", "upper_arm", "forearm", "wrist", "gripper"]
+
+        if link_num <= len(default_names):
+            default_name = default_names[link_num - 1]
+        else:
+            default_name = f"link_{link_num}"
+
+        # For now, use default name (interactive naming could be added)
+        name = default_name
+        is_fixed = (link_num == 1)  # First link is typically the base
+
+        self.links.append((name, [float(x), float(y)], is_fixed))
+        logger.info(f"Added link: {name} at ({x}, {y}), fixed={is_fixed}")
+
+    def run(self) -> List[Tuple[str, List[float], bool]]:
+        """
+        Run the interactive annotation session.
+
+        Returns:
+            List of (name, [x, y], is_fixed) tuples
+        """
+        cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL)
+        cv2.setMouseCallback(self.window_name, self._mouse_callback)
+
+        print("\n" + "=" * 60)
+        print("LINK ANNOTATION MODE")
+        print("=" * 60)
+        print("Click on each rigid part of the robot, starting with the base.")
+        print("Press [Enter] or [Space] when done.")
+        print("=" * 60 + "\n")
+
+        while self._running:
+            display = self._draw_overlay()
+            cv2.imshow(self.window_name, display)
+
+            key = cv2.waitKey(30) & 0xFF
+
+            if key == 27:  # Escape
+                print("Annotation cancelled")
+                self.links = []
+                break
+            elif key in [13, 32]:  # Enter or Space
+                if len(self.links) < 2:
+                    print("Please annotate at least 2 links (base + 1 movable)")
+                else:
+                    print(f"Annotation complete: {len(self.links)} links")
+                    break
+            elif key == ord('u') or key == ord('U'):
+                if self.links:
+                    removed = self.links.pop()
+                    print(f"Removed link: {removed[0]}")
+            elif key == ord('b') or key == ord('B'):
+                if self.links:
+                    name, point, _ = self.links[-1]
+                    self.links[-1] = (name, point, True)
+                    print(f"Marked '{name}' as base (fixed)")
+
+        cv2.destroyWindow(self.window_name)
+        return self.links
+
+
+class VideoCapture:
+    """
+    Webcam video capture with guidance overlay.
+
+    Provides real-time feedback during robot motion capture.
+    """
+
+    def __init__(self, config: CaptureConfig):
+        """
+        Initialize video capture.
+
+        Args:
+            config: Capture configuration
+        """
+        if not CV2_AVAILABLE:
+            raise CaptureError("OpenCV not installed. Run: pip install opencv-python")
+
+        self.config = config
+        self.cap: Optional[cv2.VideoCapture] = None
+        self.writer: Optional[cv2.VideoWriter] = None
+        self.frames: List["np.ndarray"] = []
+        self.recording = False
+        self.frame_count = 0
+        self.start_time: Optional[float] = None
+
+    def open(self) -> None:
+        """Open the camera."""
+        self.cap = cv2.VideoCapture(self.config.camera_id)
+
+        if not self.cap.isOpened():
+            raise CaptureError(
+                f"Could not open camera {self.config.camera_id}. "
+                "Check that camera is connected and not in use."
+            )
+
+        # Set resolution
+        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, self.config.width)
+        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, self.config.height)
+        self.cap.set(cv2.CAP_PROP_FPS, self.config.fps)
+
+        # Get actual resolution (may differ from requested)
+        actual_w = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        actual_h = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        actual_fps = self.cap.get(cv2.CAP_PROP_FPS)
+
+        logger.info(f"Camera opened: {actual_w}x{actual_h} @ {actual_fps:.1f} FPS")
+
+    def close(self) -> None:
+        """Release camera and writer resources."""
+        if self.writer:
+            self.writer.release()
+            self.writer = None
+        if self.cap:
+            self.cap.release()
+            self.cap = None
+
+    def read_frame(self) -> Optional["np.ndarray"]:
+        """Read a single frame from camera."""
+        if not self.cap:
+            return None
+
+        ret, frame = self.cap.read()
+        if not ret:
+            return None
+
+        if self.config.flip_horizontal:
+            frame = cv2.flip(frame, 1)
+
+        return frame
+
+    def _draw_guidance_overlay(
+        self,
+        frame: "np.ndarray",
+        phase: str,
+        elapsed: float,
+    ) -> "np.ndarray":
+        """Draw capture guidance overlay."""
+        display = frame.copy()
+        h, w = display.shape[:2]
+
+        # Semi-transparent panel at top
+        panel_h = 100
+        overlay = display.copy()
+        cv2.rectangle(overlay, (10, 10), (w - 10, panel_h), (0, 0, 0), -1)
+        cv2.addWeighted(overlay, 0.7, display, 0.3, 0, display)
+
+        # Status text
+        if phase == "preview":
+            status = "PREVIEW - Press [Space] to start recording"
+            color = (255, 255, 0)  # Yellow
+        elif phase == "recording":
+            status = f"RECORDING - {elapsed:.1f}s - Press [Space] to stop"
+            color = (0, 0, 255)  # Red
+
+            # Recording indicator
+            if int(elapsed * 2) % 2 == 0:
+                cv2.circle(display, (w - 40, 30), 10, (0, 0, 255), -1)
+        else:
+            status = phase
+            color = (255, 255, 255)
+
+        cv2.putText(
+            display, status, (20, 35),
+            cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2
+        )
+
+        # Guidance
+        guidance = [
+            "1. Move robot through FULL range of motion",
+            "2. Move camera around robot (or rotate robot)",
+            "   to capture from multiple angles",
+        ]
+
+        y_offset = 55
+        for line in guidance:
+            cv2.putText(
+                display, line, (20, y_offset),
+                cv2.FONT_HERSHEY_SIMPLEX, 0.45, (200, 200, 200), 1
+            )
+            y_offset += 16
+
+        # Frame counter during recording
+        if self.recording:
+            cv2.putText(
+                display, f"Frames: {self.frame_count}", (w - 150, h - 20),
+                cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1
+            )
+
+        return display
+
+    def capture_interactive(
+        self,
+        output_path: Path,
+        min_duration: float = 5.0,
+        max_duration: float = 60.0,
+    ) -> Tuple["np.ndarray", int, float]:
+        """
+        Run interactive capture session.
+
+        Args:
+            output_path: Path to save video file
+            min_duration: Minimum recording duration (seconds)
+            max_duration: Maximum recording duration (seconds)
+
+        Returns:
+            Tuple of (first_frame, frame_count, fps)
+        """
+        if not self.cap:
+            self.open()
+
+        window_name = "URDF Scan - Video Capture"
+        cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
+
+        print("\n" + "=" * 60)
+        print("VIDEO CAPTURE MODE")
+        print("=" * 60)
+        print("Press [Space] to start/stop recording")
+        print("Press [Esc] to cancel")
+        print("=" * 60 + "\n")
+
+        first_frame = None
+        phase = "preview"
+
+        fourcc = cv2.VideoWriter_fourcc(*self.config.codec)
+
+        try:
+            while True:
+                frame = self.read_frame()
+                if frame is None:
+                    raise CaptureError("Failed to read frame from camera")
+
+                if first_frame is None and self.recording:
+                    first_frame = frame.copy()
+
+                # Compute elapsed time
+                elapsed = 0.0
+                if self.start_time:
+                    elapsed = time.time() - self.start_time
+
+                # Draw overlay
+                display = self._draw_guidance_overlay(frame, phase, elapsed)
+                cv2.imshow(window_name, display)
+
+                # Handle recording
+                if self.recording:
+                    if self.writer is None:
+                        h, w = frame.shape[:2]
+                        self.writer = cv2.VideoWriter(
+                            str(output_path), fourcc,
+                            self.config.fps, (w, h)
+                        )
+                    self.writer.write(frame)
+                    self.frame_count += 1
+
+                    # Auto-stop at max duration
+                    if elapsed >= max_duration:
+                        print(f"\nMax duration reached ({max_duration}s)")
+                        break
+
+                # Handle keyboard input
+                key = cv2.waitKey(1) & 0xFF
+
+                if key == 27:  # Escape
+                    print("\nCapture cancelled")
+                    self.close()
+                    cv2.destroyWindow(window_name)
+                    raise CaptureError("Capture cancelled by user")
+
+                elif key == 32:  # Space
+                    if not self.recording:
+                        # Start recording
+                        self.recording = True
+                        self.start_time = time.time()
+                        phase = "recording"
+                        print("Recording started...")
+                    else:
+                        # Stop recording
+                        if elapsed < min_duration:
+                            print(f"Please record at least {min_duration}s "
+                                  f"(current: {elapsed:.1f}s)")
+                        else:
+                            print(f"\nRecording stopped: {self.frame_count} frames")
+                            break
+
+        finally:
+            cv2.destroyWindow(window_name)
+
+        if first_frame is None:
+            raise CaptureError("No frames captured")
+
+        # Compute actual FPS
+        if self.start_time:
+            actual_duration = time.time() - self.start_time
+            actual_fps = self.frame_count / actual_duration if actual_duration > 0 else self.config.fps
+        else:
+            actual_fps = self.config.fps
+
+        self.close()
+        return first_frame, self.frame_count, actual_fps
+
+
+def load_video(video_path: str) -> Tuple["np.ndarray", int, float, Tuple[int, int]]:
+    """
+    Load an existing video file.
+
+    Args:
+        video_path: Path to video file
+
+    Returns:
+        Tuple of (first_frame, frame_count, fps, (width, height))
+    """
+    if not CV2_AVAILABLE:
+        raise CaptureError("OpenCV not installed. Run: pip install opencv-python")
+
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        raise CaptureError(f"Could not open video: {video_path}")
+
+    # Get video properties
+    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+    # Read first frame
+    ret, first_frame = cap.read()
+    cap.release()
+
+    if not ret:
+        raise CaptureError(f"Could not read first frame from: {video_path}")
+
+    logger.info(f"Loaded video: {frame_count} frames, {width}x{height} @ {fps:.1f} FPS")
+    return first_frame, frame_count, fps, (width, height)
+
+
+def run_capture(
+    output_dir: str,
+    video_path: Optional[str] = None,
+    camera_id: int = 0,
+    robot_name: Optional[str] = None,
+    scale_ref: Optional[str] = None,
+) -> "ScanSession":
+    """
+    Run the capture phase of URDF scanning.
+
+    Args:
+        output_dir: Directory to save session data
+        video_path: Optional path to existing video (skip capture)
+        camera_id: Camera device ID for live capture
+        robot_name: Name for the robot
+        scale_ref: Scale reference (e.g., "gripper:85mm")
+
+    Returns:
+        ScanSession with captured data
+    """
+    from .scan_session import ScanSession, LinkAnnotation
+
+    # Create session
+    session = ScanSession.create(
+        output_dir=output_dir,
+        robot_name=robot_name,
+        scale_ref=scale_ref,
+    )
+
+    # Get video (capture or load)
+    if video_path:
+        # Load existing video
+        first_frame, frame_count, fps, (width, height) = load_video(video_path)
+
+        # Copy video to session directory
+        import shutil
+        shutil.copy(video_path, session.video_path)
+        print(f"Loaded video: {video_path}")
+
+    else:
+        # Interactive capture
+        config = CaptureConfig(camera_id=camera_id)
+        capture = VideoCapture(config)
+        capture.open()
+
+        first_frame, frame_count, fps = capture.capture_interactive(
+            session.video_path,
+            min_duration=5.0,
+            max_duration=60.0,
+        )
+        height, width = first_frame.shape[:2]
+
+    # Update metadata
+    session.metadata.frame_count = frame_count
+    session.metadata.fps = fps
+    session.metadata.resolution = [width, height]
+    session.metadata.video_path = str(session.video_path)
+
+    # Run link annotation
+    print("\nStarting link annotation...")
+    annotator = LinkAnnotator(first_frame)
+    link_data = annotator.run()
+
+    if not link_data:
+        raise CaptureError("No links annotated. Capture incomplete.")
+
+    # Convert to LinkAnnotation objects
+    session.links = [
+        LinkAnnotation(name=name, point=point, is_fixed=is_fixed)
+        for name, point, is_fixed in link_data
+    ]
+
+    # Save session data
+    session.save_links()
+    session.metadata.capture_complete = True
+    session.save_metadata()
+
+    print(f"\nCapture complete!")
+    print(f"  Video: {session.video_path}")
+    print(f"  Frames: {frame_count}")
+    print(f"  Links: {len(session.links)}")
+    for link in session.links:
+        marker = " (base)" if link.is_fixed else ""
+        print(f"    - {link.name}{marker} at {link.point}")
+
+    return session
+
+
+__all__ = [
+    "CaptureError",
+    "CaptureConfig",
+    "LinkAnnotator",
+    "VideoCapture",
+    "load_video",
+    "run_capture",
+]