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

ate/robot/teach.py

@@ -0,0 +1,515 @@
+"""
+Fast Teaching Mode for Robot Skill Development.
+
+Provides a real-time, keyboard-driven interface for:
+1. Live webcam preview
+2. Quick servo adjustment with visual feedback
+3. One-key pose saving
+4. Immediate playback testing
+
+Usage:
+    ate robot teach --port /dev/cu.usbserial-10 --name mechdog
+
+Controls:
+    ↑/↓     Select servo
+    ←/→     Decrease/increase servo value
+    SHIFT   Fine adjustment (±10 instead of ±50)
+    SPACE   Move to current position
+    S       Save current position as pose
+    A       Create action from saved poses
+    P       Playback last action
+    R       Reset to center positions
+    Q       Quit and save
+"""
+
+import sys
+import time
+import json
+import threading
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+from dataclasses import dataclass, field
+from datetime import datetime
+
+try:
+    import cv2
+    HAS_OPENCV = True
+except ImportError:
+    HAS_OPENCV = False
+
+from .calibration import (
+    RobotCalibration,
+    ServoCalibration,
+    Pose,
+    JointType,
+    load_calibration,
+    save_calibration,
+)
+from .visual_labeler import (
+    SkillLibrary,
+    Action,
+    ActionStep,
+    ActionType,
+    load_skill_library,
+    save_skill_library,
+    ROBOT_PRESETS,
+    apply_preset_to_calibration,
+)
+
+
+@dataclass
+class TeachingSession:
+    """State for a teaching session."""
+    robot_name: str
+    calibration: RobotCalibration
+    library: SkillLibrary
+
+    # Current state
+    selected_servo_idx: int = 0
+    servo_ids: List[int] = field(default_factory=list)
+    current_positions: Dict[int, int] = field(default_factory=dict)
+
+    # Recorded poses for action creation
+    recorded_poses: List[str] = field(default_factory=list)
+
+    # Display
+    window_name: str = "Robot Teaching Mode"
+
+    def __post_init__(self):
+        self.servo_ids = sorted(self.calibration.servos.keys())
+        # Initialize current positions to center
+        for sid in self.servo_ids:
+            servo = self.calibration.servos[sid]
+            self.current_positions[sid] = servo.center_value
+
+
+class TeachingInterface:
+    """
+    Real-time teaching interface with live preview.
+
+    Uses OpenCV for display and keyboard handling.
+    """
+
+    def __init__(
+        self,
+        serial_port: str,
+        robot_name: str = "robot",
+        robot_model: str = "unknown",
+        webcam_id: int = 0,
+        camera_url: Optional[str] = None,
+    ):
+        if not HAS_OPENCV:
+            raise RuntimeError("OpenCV required. Install with: pip install opencv-python")
+
+        self.serial_port = serial_port
+        self.robot_name = robot_name
+        self.robot_model = robot_model
+        self.webcam_id = webcam_id
+        self.camera_url = camera_url
+
+        # Serial connection
+        self.serial = None
+        self._connect_serial()
+
+        # Camera
+        self.webcam = None
+        self.frame = None
+        self.frame_lock = threading.Lock()
+
+        # Load or create calibration
+        self.calibration = load_calibration(robot_name)
+        if not self.calibration:
+            self.calibration = RobotCalibration(
+                robot_model=robot_model,
+                robot_name=robot_name,
+                serial_port=serial_port,
+            )
+            # Apply preset if available
+            preset = ROBOT_PRESETS.get(robot_model)
+            if preset:
+                apply_preset_to_calibration(self.calibration, preset)
+                print(f"Applied {preset['name']} preset")
+
+        # Load or create skill library
+        self.library = load_skill_library(robot_name)
+        if not self.library:
+            self.library = SkillLibrary(
+                robot_name=robot_name,
+                robot_model=robot_model,
+            )
+
+        # Create session
+        self.session = TeachingSession(
+            robot_name=robot_name,
+            calibration=self.calibration,
+            library=self.library,
+        )
+
+        # Discover servos if none defined
+        if not self.session.servo_ids:
+            self._discover_servos()
+
+    def _connect_serial(self):
+        """Connect to robot via serial."""
+        try:
+            import serial
+            self.serial = serial.Serial(self.serial_port, 115200, timeout=1)
+            time.sleep(0.5)  # Wait for connection
+            # Clear any pending data
+            self.serial.reset_input_buffer()
+        except Exception as e:
+            print(f"Serial connection failed: {e}")
+            self.serial = None
+
+    def _send_command(self, cmd: str) -> Optional[str]:
+        """Send command to robot and get response."""
+        if not self.serial:
+            return None
+        try:
+            self.serial.write(f"{cmd}\n".encode())
+            self.serial.flush()
+            time.sleep(0.05)
+            response = ""
+            while self.serial.in_waiting:
+                response += self.serial.read(self.serial.in_waiting).decode('utf-8', errors='ignore')
+                time.sleep(0.01)
+            return response.strip()
+        except Exception as e:
+            print(f"Command failed: {e}")
+            return None
+
+    def _discover_servos(self):
+        """Discover connected servos."""
+        print("Discovering servos...")
+        found = []
+        for sid in range(1, 16):
+            response = self._send_command(f"servo_read({sid})")
+            if response and "error" not in response.lower():
+                try:
+                    # Try to parse position
+                    pos = int(response.split("=")[-1].strip().rstrip(")"))
+                    found.append(sid)
+                    self.session.current_positions[sid] = pos
+
+                    # Add to calibration if not present
+                    if sid not in self.calibration.servos:
+                        self.calibration.servos[sid] = ServoCalibration(
+                            servo_id=sid,
+                            name=f"servo_{sid}",
+                            joint_type=JointType.UNKNOWN,
+                        )
+                except:
+                    pass
+
+        self.session.servo_ids = sorted(found)
+        print(f"Found {len(found)} servos: {found}")
+
+    def _move_servo(self, servo_id: int, position: int, time_ms: int = 200):
+        """Move a servo to position."""
+        # Clamp to limits
+        servo = self.calibration.servos.get(servo_id)
+        if servo:
+            position = max(servo.min_value, min(servo.max_value, position))
+        else:
+            position = max(200, min(2500, position))
+
+        self._send_command(f"servo_move({servo_id}, {position}, {time_ms})")
+        self.session.current_positions[servo_id] = position
+        return position
+
+    def _start_webcam(self):
+        """Start webcam capture in background thread."""
+        def capture_loop():
+            cap = cv2.VideoCapture(self.webcam_id)
+            cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
+            cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
+
+            while self.webcam is not None:
+                ret, frame = cap.read()
+                if ret:
+                    with self.frame_lock:
+                        self.frame = frame
+                time.sleep(0.033)  # ~30 fps
+
+            cap.release()
+
+        self.webcam = threading.Thread(target=capture_loop, daemon=True)
+        self.webcam.start()
+
+    def _get_frame(self) -> Optional[any]:
+        """Get current webcam frame."""
+        with self.frame_lock:
+            return self.frame.copy() if self.frame is not None else None
+
+    def _draw_overlay(self, frame) -> any:
+        """Draw servo info overlay on frame."""
+        if frame is None:
+            # Create blank frame
+            frame = 255 * (cv2.imread.__module__ and 1)  # placeholder
+            return frame
+
+        h, w = frame.shape[:2]
+
+        # Semi-transparent overlay on right side
+        overlay = frame.copy()
+        cv2.rectangle(overlay, (w - 250, 0), (w, h), (40, 40, 40), -1)
+        frame = cv2.addWeighted(overlay, 0.7, frame, 0.3, 0)
+
+        # Title
+        cv2.putText(frame, "TEACH MODE", (w - 240, 30),
+                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)
+
+        # Servo list
+        y = 70
+        for i, sid in enumerate(self.session.servo_ids):
+            servo = self.calibration.servos.get(sid)
+            name = servo.name if servo else f"servo_{sid}"
+            pos = self.session.current_positions.get(sid, 1500)
+
+            # Highlight selected
+            is_selected = (i == self.session.selected_servo_idx)
+            color = (0, 255, 0) if is_selected else (200, 200, 200)
+            marker = ">" if is_selected else " "
+
+            # Draw servo info
+            cv2.putText(frame, f"{marker} {name[:12]}", (w - 240, y),
+                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
+
+            # Draw position bar
+            if servo:
+                pct = (pos - servo.min_value) / max(1, servo.max_value - servo.min_value)
+            else:
+                pct = (pos - 500) / 2000
+            bar_w = int(80 * pct)
+            cv2.rectangle(frame, (w - 100, y - 10), (w - 20, y), (60, 60, 60), -1)
+            cv2.rectangle(frame, (w - 100, y - 10), (w - 100 + bar_w, y), color, -1)
+            cv2.putText(frame, str(pos), (w - 100, y + 15),
+                       cv2.FONT_HERSHEY_SIMPLEX, 0.4, color, 1)
+
+            y += 45
+
+        # Controls hint
+        cv2.putText(frame, "[Arrows] Move", (w - 240, h - 100),
+                   cv2.FONT_HERSHEY_SIMPLEX, 0.4, (150, 150, 150), 1)
+        cv2.putText(frame, "[S] Save pose", (w - 240, h - 80),
+                   cv2.FONT_HERSHEY_SIMPLEX, 0.4, (150, 150, 150), 1)
+        cv2.putText(frame, "[P] Playback", (w - 240, h - 60),
+                   cv2.FONT_HERSHEY_SIMPLEX, 0.4, (150, 150, 150), 1)
+        cv2.putText(frame, "[Q] Quit", (w - 240, h - 40),
+                   cv2.FONT_HERSHEY_SIMPLEX, 0.4, (150, 150, 150), 1)
+
+        # Recorded poses
+        if self.session.recorded_poses:
+            cv2.putText(frame, f"Recorded: {len(self.session.recorded_poses)}", (10, h - 20),
+                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 1)
+
+        return frame
+
+    def _save_pose(self, name: Optional[str] = None):
+        """Save current position as a named pose."""
+        if name is None:
+            # Generate name
+            existing = len(self.calibration.poses)
+            name = f"pose_{existing + 1}"
+
+        pose = Pose(
+            name=name,
+            description="",
+            servo_positions=dict(self.session.current_positions),
+            transition_time_ms=500,
+        )
+        self.calibration.poses[name] = pose
+        self.session.recorded_poses.append(name)
+
+        # Save calibration
+        save_calibration(self.calibration)
+
+        print(f"Saved pose: {name}")
+        return name
+
+    def _apply_pose(self, pose_name: str, time_ms: int = 500):
+        """Apply a saved pose."""
+        pose = self.calibration.poses.get(pose_name)
+        if not pose:
+            print(f"Pose not found: {pose_name}")
+            return
+
+        for sid, pos in pose.servo_positions.items():
+            self._move_servo(sid, pos, time_ms)
+
+        time.sleep(time_ms / 1000 + 0.1)
+
+    def _create_action(self, name: str) -> Optional[Action]:
+        """Create an action from recorded poses."""
+        if not self.session.recorded_poses:
+            print("No poses recorded")
+            return None
+
+        steps = []
+        for pose_name in self.session.recorded_poses:
+            steps.append(ActionStep(
+                pose_name=pose_name,
+                duration_ms=500,
+                wait_after_ms=200,
+            ))
+
+        action = Action(
+            name=name,
+            action_type=ActionType.MANIPULATION,
+            description=f"Action with {len(steps)} steps",
+            steps=steps,
+        )
+
+        self.library.actions[name] = action
+        save_skill_library(self.library)
+
+        print(f"Created action: {name} ({len(steps)} steps)")
+        return action
+
+    def _playback_action(self, action_name: str):
+        """Playback an action."""
+        action = self.library.actions.get(action_name)
+        if not action:
+            print(f"Action not found: {action_name}")
+            return
+
+        print(f"Playing: {action_name}")
+        for step in action.steps:
+            self._apply_pose(step.pose_name, step.duration_ms)
+            if step.wait_after_ms > 0:
+                time.sleep(step.wait_after_ms / 1000)
+
+    def run(self):
+        """Run the teaching interface."""
+        print("\n=== TEACHING MODE ===")
+        print("Controls:")
+        print("  ↑/↓  Select servo")
+        print("  ←/→  Adjust position")
+        print("  S    Save pose")
+        print("  A    Create action from poses")
+        print("  P    Playback last action")
+        print("  C    Clear recorded poses")
+        print("  Q    Quit and save")
+        print()
+
+        # Start webcam
+        self._start_webcam()
+        time.sleep(0.5)
+
+        # Create window
+        cv2.namedWindow(self.session.window_name, cv2.WINDOW_NORMAL)
+        cv2.resizeWindow(self.session.window_name, 900, 600)
+
+        step_size = 50  # Default step size
+        last_action = None
+
+        try:
+            while True:
+                # Get and display frame
+                frame = self._get_frame()
+                if frame is not None:
+                    display = self._draw_overlay(frame)
+                    cv2.imshow(self.session.window_name, display)
+
+                # Handle keyboard
+                key = cv2.waitKey(30) & 0xFF
+
+                if key == ord('q') or key == 27:  # Q or ESC
+                    break
+
+                elif key == 82 or key == 0:  # Up arrow
+                    self.session.selected_servo_idx = max(0, self.session.selected_servo_idx - 1)
+
+                elif key == 84 or key == 1:  # Down arrow
+                    self.session.selected_servo_idx = min(
+                        len(self.session.servo_ids) - 1,
+                        self.session.selected_servo_idx + 1
+                    )
+
+                elif key == 81 or key == 2:  # Left arrow
+                    if self.session.servo_ids:
+                        sid = self.session.servo_ids[self.session.selected_servo_idx]
+                        pos = self.session.current_positions.get(sid, 1500)
+                        self._move_servo(sid, pos - step_size)
+
+                elif key == 83 or key == 3:  # Right arrow
+                    if self.session.servo_ids:
+                        sid = self.session.servo_ids[self.session.selected_servo_idx]
+                        pos = self.session.current_positions.get(sid, 1500)
+                        self._move_servo(sid, pos + step_size)
+
+                elif key == ord('s'):  # Save pose
+                    name = self._save_pose()
+                    print(f"Pose saved: {name}")
+
+                elif key == ord('a'):  # Create action
+                    if self.session.recorded_poses:
+                        action_name = f"action_{len(self.library.actions) + 1}"
+                        self._create_action(action_name)
+                        last_action = action_name
+                        self.session.recorded_poses = []  # Clear for next action
+                    else:
+                        print("No poses recorded. Press S to save poses first.")
+
+                elif key == ord('p'):  # Playback
+                    if last_action:
+                        self._playback_action(last_action)
+                    elif self.library.actions:
+                        last_action = list(self.library.actions.keys())[-1]
+                        self._playback_action(last_action)
+                    else:
+                        print("No actions to play")
+
+                elif key == ord('c'):  # Clear recorded
+                    self.session.recorded_poses = []
+                    print("Cleared recorded poses")
+
+                elif key == ord('r'):  # Reset to center
+                    for sid in self.session.servo_ids:
+                        servo = self.calibration.servos.get(sid)
+                        center = servo.center_value if servo else 1500
+                        self._move_servo(sid, center)
+                    print("Reset to center positions")
+
+                elif key == ord('+') or key == ord('='):
+                    step_size = min(200, step_size + 10)
+                    print(f"Step size: {step_size}")
+
+                elif key == ord('-'):
+                    step_size = max(10, step_size - 10)
+                    print(f"Step size: {step_size}")
+
+        finally:
+            # Cleanup
+            self.webcam = None  # Stop capture thread
+            cv2.destroyAllWindows()
+
+            if self.serial:
+                self.serial.close()
+
+            # Final save
+            save_calibration(self.calibration)
+            save_skill_library(self.library)
+            print("\nSession saved.")
+
+
+def teach_command(
+    port: str,
+    name: str = "robot",
+    robot_type: str = "hiwonder_mechdog",
+    webcam_id: int = 0,
+    camera_url: Optional[str] = None,
+):
+    """
+    Run the teaching interface.
+
+    ate robot teach --port /dev/cu.usbserial-10 --name mechdog
+    """
+    interface = TeachingInterface(
+        serial_port=port,
+        robot_name=name,
+        robot_model=robot_type,
+        webcam_id=webcam_id,
+        camera_url=camera_url,
+    )
+    interface.run()