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

ate/recording/visual.py

@@ -0,0 +1,416 @@
+"""
+Visual recording for demonstrations with image capture and detection.
+
+Extends the standard recording session to capture camera images at regular
+intervals, optionally running object detection on each frame.
+
+This enables:
+- Recording demonstrations with synchronized video
+- Automatic object detection during recording
+- Training data collection for perception models
+"""
+
+import time
+import json
+import base64
+import threading
+from pathlib import Path
+from dataclasses import dataclass, field
+from typing import List, Optional, Dict, Any, Callable
+
+from ate.interfaces import RobotInterface, CameraInterface
+from ate.interfaces.perception import Image
+from .session import RecordingSession, RecordedCall, RecordingMetadata
+
+
+@dataclass
+class DetectionResult:
+    """A single object detection result."""
+    label: str
+    confidence: float
+    bbox: Dict[str, int]  # x, y, width, height
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> dict:
+        return {
+            "label": self.label,
+            "confidence": self.confidence,
+            "bbox": self.bbox,
+            "metadata": self.metadata,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "DetectionResult":
+        return cls(**data)
+
+
+@dataclass
+class RecordedFrame:
+    """A recorded camera frame with optional detections."""
+
+    timestamp: float  # Unix timestamp
+    relative_time: float  # Time since recording started
+    width: int
+    height: int
+    format: str  # jpeg, png, rgb8
+    data: bytes  # Raw image data
+    detections: List[DetectionResult] = field(default_factory=list)
+
+    def to_dict(self, include_data: bool = True) -> dict:
+        result = {
+            "timestamp": self.timestamp,
+            "relative_time": self.relative_time,
+            "width": self.width,
+            "height": self.height,
+            "format": self.format,
+            "detections": [d.to_dict() for d in self.detections],
+        }
+        if include_data:
+            result["data"] = base64.b64encode(self.data).decode('ascii')
+        return result
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "RecordedFrame":
+        frame_data = base64.b64decode(data.get("data", ""))
+        detections = [DetectionResult.from_dict(d) for d in data.get("detections", [])]
+        return cls(
+            timestamp=data["timestamp"],
+            relative_time=data["relative_time"],
+            width=data["width"],
+            height=data["height"],
+            format=data["format"],
+            data=frame_data,
+            detections=detections,
+        )
+
+
+class VisualRecordingSession(RecordingSession):
+    """
+    Recording session with visual capture and detection.
+
+    Extends RecordingSession to capture camera images at regular intervals
+    and optionally run object detection on each frame.
+
+    Usage:
+        from ate.detection import TrashDetector
+
+        detector = TrashDetector()
+
+        with VisualRecordingSession(
+            driver,
+            name="pickup_demo",
+            capture_fps=5.0,
+            detector=detector,
+        ) as session:
+            # Drive robot - images captured automatically
+            driver.walk(Vector3.forward(), speed=0.2)
+            time.sleep(5.0)
+
+        # Save with visual data
+        session.save("pickup_demo.demonstration")
+        session.save_frames("pickup_demo_frames.json")
+    """
+
+    def __init__(
+        self,
+        driver: RobotInterface,
+        name: str = "recording",
+        description: Optional[str] = None,
+        tags: Optional[List[str]] = None,
+        auto_save: bool = False,
+        save_path: Optional[str] = None,
+        # Visual capture settings
+        capture_fps: float = 5.0,  # Frames per second to capture
+        detector: Any = None,  # Object detector (TrashDetector, ColorDetector, etc.)
+        max_frames: int = 1000,  # Maximum frames to store
+        store_frames: bool = True,  # Whether to store image data
+    ):
+        """
+        Initialize visual recording session.
+
+        Args:
+            driver: Robot driver (must implement CameraInterface for visual capture)
+            name: Human-readable name
+            description: Optional description
+            tags: Optional tags
+            auto_save: Auto-save on exit
+            save_path: Path for auto-save
+            capture_fps: Frames per second to capture (default: 5.0)
+            detector: Optional object detector for automatic detection
+            max_frames: Maximum frames to store in memory
+            store_frames: Whether to store raw image data
+        """
+        super().__init__(driver, name, description, tags, auto_save, save_path)
+
+        self._capture_fps = capture_fps
+        self._detector = detector
+        self._max_frames = max_frames
+        self._store_frames = store_frames
+
+        self._frames: List[RecordedFrame] = []
+        self._capture_thread: Optional[threading.Thread] = None
+        self._stop_capture = threading.Event()
+
+        # Check if driver has camera
+        self._has_camera = isinstance(driver, CameraInterface)
+
+    @property
+    def frames(self) -> List[RecordedFrame]:
+        """Get all recorded frames."""
+        return self._frames.copy()
+
+    @property
+    def frame_count(self) -> int:
+        """Get number of captured frames."""
+        return len(self._frames)
+
+    @property
+    def detection_count(self) -> int:
+        """Get total number of detections across all frames."""
+        return sum(len(f.detections) for f in self._frames)
+
+    def start(self) -> None:
+        """Start recording with visual capture."""
+        super().start()
+
+        if self._has_camera:
+            self._stop_capture.clear()
+            self._capture_thread = threading.Thread(
+                target=self._capture_loop,
+                daemon=True,
+            )
+            self._capture_thread.start()
+
+    def stop(self) -> None:
+        """Stop recording and visual capture."""
+        self._stop_capture.set()
+
+        if self._capture_thread and self._capture_thread.is_alive():
+            self._capture_thread.join(timeout=2.0)
+
+        super().stop()
+
+    def _capture_loop(self) -> None:
+        """Background thread for capturing frames."""
+        interval = 1.0 / self._capture_fps
+
+        while not self._stop_capture.is_set():
+            if len(self._frames) >= self._max_frames:
+                break
+
+            try:
+                self._capture_frame()
+            except Exception as e:
+                # Log but don't crash on capture errors
+                pass
+
+            self._stop_capture.wait(interval)
+
+    def _capture_frame(self) -> None:
+        """Capture a single frame and run detection."""
+        if not self._has_camera:
+            return
+
+        timestamp = time.time()
+        relative_time = timestamp - self._start_time
+
+        # Get image from camera
+        image = self._driver.get_image()
+        if image.width == 0 or image.height == 0:
+            return  # Skip empty frames
+
+        # Run detection if detector available
+        detections = []
+        if self._detector is not None:
+            try:
+                # Try detect_trash first (for TrashDetector)
+                if hasattr(self._detector, 'detect_trash'):
+                    items = self._detector.detect_trash(image)
+                    for item in items:
+                        det = item.detection
+                        detections.append(DetectionResult(
+                            label=item.trash_type,
+                            confidence=item.confidence,
+                            bbox={
+                                "x": det.bbox.x,
+                                "y": det.bbox.y,
+                                "width": det.bbox.width,
+                                "height": det.bbox.height,
+                            },
+                            metadata={
+                                "is_on_ground": item.is_on_ground,
+                                "size_category": item.size_category,
+                                "priority": item.priority,
+                            },
+                        ))
+                else:
+                    # Generic detector with detect() method
+                    raw_detections = self._detector.detect(image)
+                    for det in raw_detections:
+                        detections.append(DetectionResult(
+                            label=det.label,
+                            confidence=det.confidence,
+                            bbox={
+                                "x": det.bbox.x,
+                                "y": det.bbox.y,
+                                "width": det.bbox.width,
+                                "height": det.bbox.height,
+                            },
+                            metadata=det.metadata,
+                        ))
+            except Exception as e:
+                # Detection failed - still record frame without detections
+                pass
+
+        # Create recorded frame
+        frame = RecordedFrame(
+            timestamp=timestamp,
+            relative_time=relative_time,
+            width=image.width,
+            height=image.height,
+            format=image.encoding,
+            data=image.data if self._store_frames else b"",
+            detections=detections,
+        )
+
+        self._frames.append(frame)
+
+    def capture_now(self) -> Optional[RecordedFrame]:
+        """
+        Manually capture a frame right now.
+
+        Useful for capturing at specific moments (e.g., before/after actions).
+
+        Returns:
+            Captured frame or None if capture failed
+        """
+        if not self._has_camera or not self._recording:
+            return None
+
+        prev_count = len(self._frames)
+        self._capture_frame()
+
+        if len(self._frames) > prev_count:
+            return self._frames[-1]
+        return None
+
+    def get_frames_in_range(
+        self,
+        start_time: float,
+        end_time: float,
+    ) -> List[RecordedFrame]:
+        """Get frames within a time range."""
+        return [
+            f for f in self._frames
+            if start_time <= f.relative_time <= end_time
+        ]
+
+    def get_frames_with_detections(
+        self,
+        label: Optional[str] = None,
+        min_confidence: float = 0.0,
+    ) -> List[RecordedFrame]:
+        """Get frames that have detections matching criteria."""
+        result = []
+        for frame in self._frames:
+            matching = [
+                d for d in frame.detections
+                if d.confidence >= min_confidence
+                and (label is None or d.label == label)
+            ]
+            if matching:
+                result.append(frame)
+        return result
+
+    def get_metadata(self) -> RecordingMetadata:
+        """Get recording metadata including visual stats."""
+        metadata = super().get_metadata()
+        # Add visual metadata to tags
+        if self._has_camera:
+            metadata.tags = list(metadata.tags) + [
+                f"frames:{len(self._frames)}",
+                f"detections:{self.detection_count}",
+            ]
+        return metadata
+
+    def save_frames(self, path: str, include_data: bool = True) -> None:
+        """
+        Save frames to a separate file.
+
+        Args:
+            path: File path for frames
+            include_data: Whether to include raw image data (base64)
+        """
+        data = {
+            "version": "1.0",
+            "recording_id": self.get_metadata().id,
+            "frame_count": len(self._frames),
+            "detection_count": self.detection_count,
+            "frames": [f.to_dict(include_data=include_data) for f in self._frames],
+        }
+
+        with open(path, 'w') as f:
+            json.dump(data, f, indent=2 if not include_data else None)
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary including frame summary."""
+        base = super().to_dict()
+
+        # Add visual summary (not full frames - too large for API)
+        base["visual"] = {
+            "has_camera": self._has_camera,
+            "frame_count": len(self._frames),
+            "detection_count": self.detection_count,
+            "capture_fps": self._capture_fps,
+            "detection_summary": self._get_detection_summary(),
+        }
+
+        return base
+
+    def _get_detection_summary(self) -> Dict[str, int]:
+        """Get count of each detection type."""
+        counts: Dict[str, int] = {}
+        for frame in self._frames:
+            for det in frame.detections:
+                counts[det.label] = counts.get(det.label, 0) + 1
+        return counts
+
+    def summary(self) -> str:
+        """Get human-readable summary including visual data."""
+        base = super().summary()
+
+        if not self._has_camera:
+            return base + "\n\nVisual capture: Not available (no camera)"
+
+        lines = [
+            base,
+            "",
+            "Visual capture:",
+            f"  Frames captured: {len(self._frames)}",
+            f"  Total detections: {self.detection_count}",
+        ]
+
+        # Detection breakdown
+        summary = self._get_detection_summary()
+        if summary:
+            lines.append("  Detection breakdown:")
+            for label, count in sorted(summary.items(), key=lambda x: -x[1]):
+                lines.append(f"    - {label}: {count}")
+
+        return "\n".join(lines)
+
+
+def load_frames(path: str) -> List[RecordedFrame]:
+    """
+    Load recorded frames from a file.
+
+    Args:
+        path: Path to frames file
+
+    Returns:
+        List of RecordedFrame objects
+    """
+    with open(path, 'r') as f:
+        data = json.load(f)
+
+    return [RecordedFrame.from_dict(frame) for frame in data.get("frames", [])]

ate/recording/wrapper.py

@@ -0,0 +1,95 @@
+"""
+Recording wrapper for automatic telemetry capture.
+
+Provides a simpler API than RecordingSession for cases where you
+want to wrap a driver and record everything automatically.
+"""
+
+from typing import Optional
+from ate.interfaces import RobotInterface
+from .session import RecordingSession
+
+
+class RecordingWrapper:
+    """
+    Wrapper that automatically records all driver interactions.
+
+    Unlike RecordingSession (context manager), RecordingWrapper creates
+    a persistent recording that captures everything until explicitly stopped.
+
+    Usage:
+        # Wrap a driver
+        raw_driver = MechDogDriver(port="/dev/...")
+        driver = RecordingWrapper(raw_driver, name="session_01")
+
+        # Use driver normally - all calls are recorded
+        driver.connect()
+        driver.stand()
+        driver.walk(Vector3.forward(), speed=0.3)
+        driver.stop()
+        driver.disconnect()
+
+        # Save the recording
+        driver.save_recording("session_01.demonstration")
+    """
+
+    def __init__(
+        self,
+        driver: RobotInterface,
+        name: str = "recording",
+        description: Optional[str] = None,
+        auto_start: bool = True,
+    ):
+        """
+        Initialize recording wrapper.
+
+        Args:
+            driver: Robot driver to wrap
+            name: Recording name
+            description: Optional description
+            auto_start: Start recording immediately
+        """
+        self._driver = driver
+        self._session = RecordingSession(
+            driver,
+            name=name,
+            description=description,
+        )
+
+        if auto_start:
+            self._session.start()
+
+    def __getattr__(self, name: str):
+        """Forward attribute access to wrapped driver."""
+        return getattr(self._driver, name)
+
+    def start_recording(self) -> None:
+        """Start recording (if not already)."""
+        self._session.start()
+
+    def stop_recording(self) -> None:
+        """Stop recording."""
+        self._session.stop()
+
+    def save_recording(self, path: str) -> None:
+        """Save recording to file."""
+        self._session.save(path)
+
+    def get_recording_summary(self) -> str:
+        """Get recording summary."""
+        return self._session.summary()
+
+    @property
+    def is_recording(self) -> bool:
+        """Check if currently recording."""
+        return self._session.is_recording
+
+    @property
+    def recording_duration(self) -> Optional[float]:
+        """Get recording duration in seconds."""
+        return self._session.duration
+
+    @property
+    def call_count(self) -> int:
+        """Get number of recorded calls."""
+        return len(self._session.calls)