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

ate/robot/direction_calibration.py

@@ -0,0 +1,554 @@
+"""
+Direction Calibration via Twitch Test.
+
+This module implements the CRITICAL missing step in robot calibration:
+determining whether positive servo commands move TOWARD or AWAY from targets.
+
+The Problem:
+-----------
+Previous calibration discovered:
+- Which servos exist and their ranges
+- Which servo controls which joint (kinematic structure)
+
+But NOT:
+- Does servo 8 at 1700 move gripper TOWARD the ball or AWAY from it?
+
+Without this, behaviors like pickup use wrong directions and fail catastrophically.
+
+The Solution - Twitch Test:
+--------------------------
+1. Place target (ball) in front of robot
+2. Detect gripper position via ArUco marker
+3. For each servo:
+   a. Measure distance(gripper, ball) BEFORE
+   b. Move servo by +100 (small perturbation)
+   c. Measure distance(gripper, ball) AFTER
+   d. If distance decreased → positive = toward target
+   e. If distance increased → positive = away from target
+4. Store direction mapping
+5. Behaviors now use correct signs!
+
+Usage:
+    from ate.robot.direction_calibration import run_direction_calibration
+
+    # Run interactive calibration
+    result = run_direction_calibration(
+        serial_port="/dev/cu.usbserial-10",
+        robot_name="my_mechdog",
+    )
+
+    # Result contains direction mappings for all arm servos
+"""
+
+import cv2
+import numpy as np
+import serial
+import time
+import json
+import sys
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+from typing import Optional, Dict, List, Tuple, Any
+
+from .calibration_state import (
+    CalibrationState,
+    DirectionMapping,
+    require_direction_calibration,
+)
+
+
+@dataclass
+class TwitchResult:
+    """Result of a single twitch test for one servo."""
+    servo_id: int
+    initial_distance: float  # pixels
+    final_distance: float    # pixels
+    delta: float            # final - initial (negative = got closer)
+    perturbation: int       # PWM units moved
+    direction: str          # "toward" or "away"
+    confidence: float       # based on magnitude of delta
+
+
+class DirectionCalibrator:
+    """
+    Automated direction calibration using ArUco + ball detection.
+
+    The POINT of ArUco markers + webcam:
+    - Camera sees ball position
+    - Camera sees gripper marker
+    - Move servo → measure if gripper got closer or further from ball
+    - NO HUMAN OBSERVATION NEEDED
+    """
+
+    def __init__(
+        self,
+        serial_port: str,
+        robot_name: str = "robot",
+        gripper_marker_id: int = 15,
+        camera_index: int = 0,
+        arm_servos: Optional[List[int]] = None,
+    ):
+        self.serial_port = serial_port
+        self.robot_name = robot_name
+        self.gripper_marker_id = gripper_marker_id
+        self.camera_index = camera_index
+        self.arm_servos = arm_servos or [8, 9, 10, 11]  # Default MechDog arm servos
+
+        self.cap: Optional[cv2.VideoCapture] = None
+        self.ser: Optional[serial.Serial] = None
+        self.aruco_detector = None
+
+        # Results
+        self.results: Dict[int, TwitchResult] = {}
+
+    def setup_camera(self) -> bool:
+        """Initialize webcam and ArUco detector."""
+        print("Initializing camera...")
+
+        self.cap = cv2.VideoCapture(self.camera_index)
+        if not self.cap.isOpened():
+            print(f"Error: Could not open camera {self.camera_index}")
+            return False
+
+        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
+        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
+
+        # Warm up camera
+        for _ in range(10):
+            self.cap.read()
+
+        # Setup ArUco detector
+        aruco_dict = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_50)
+        aruco_params = cv2.aruco.DetectorParameters()
+        self.aruco_detector = cv2.aruco.ArucoDetector(aruco_dict, aruco_params)
+
+        print("Camera initialized.")
+        return True
+
+    def setup_robot(self) -> bool:
+        """Connect to robot via serial."""
+        print(f"Connecting to robot on {self.serial_port}...")
+
+        try:
+            self.ser = serial.Serial(self.serial_port, 115200, timeout=2)
+            time.sleep(0.5)
+            self.ser.read(self.ser.in_waiting)  # Clear buffer
+
+            # Initialize MechDog
+            self.send_command("from HW_MechDog import MechDog", wait=1.5)
+            self.send_command("dog = MechDog()", wait=1.5)
+
+            print("Robot connected.")
+            return True
+
+        except Exception as e:
+            print(f"Error connecting to robot: {e}")
+            return False
+
+    def send_command(self, command: str, wait: float = 0.5) -> str:
+        """Send command to robot."""
+        if not self.ser:
+            return ""
+
+        try:
+            self.ser.write(f"{command}\r\n".encode())
+            time.sleep(wait)
+            return self.ser.read(self.ser.in_waiting).decode("utf-8", errors="ignore")
+        except Exception as e:
+            print(f"Command error: {e}")
+            return ""
+
+    def set_servo(self, servo_id: int, position: int, duration: int = 500) -> None:
+        """Set a servo to a specific position."""
+        self.send_command(f"dog.set_servo({servo_id}, {position}, {duration})", wait=duration/1000 + 0.3)
+
+    def detect_ball(self, frame: np.ndarray) -> Optional[Tuple[int, int]]:
+        """Detect green/cyan ball, return center (x, y) or None."""
+        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
+
+        # Blue-green ball detection (handles green, cyan, teal balls)
+        # Expanded range to handle different lighting conditions
+        lower = np.array([70, 50, 50])
+        upper = np.array([110, 255, 255])
+        mask = cv2.inRange(hsv, lower, upper)
+
+        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+
+        if contours:
+            largest = max(contours, key=cv2.contourArea)
+            if cv2.contourArea(largest) > 500:
+                M = cv2.moments(largest)
+                if M["m00"] > 0:
+                    cx = int(M["m10"] / M["m00"])
+                    cy = int(M["m01"] / M["m00"])
+                    return (cx, cy)
+
+        return None
+
+    def detect_markers(self, frame: np.ndarray) -> Dict[int, Tuple[float, float]]:
+        """Detect ArUco markers, return dict of marker_id -> (x, y)."""
+        if self.aruco_detector is None:
+            return {}
+
+        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+        corners, ids, _ = self.aruco_detector.detectMarkers(gray)
+
+        markers = {}
+        if ids is not None:
+            for i, mid in enumerate(ids.flatten()):
+                c = corners[i][0]
+                center = (float(np.mean(c[:, 0])), float(np.mean(c[:, 1])))
+                markers[int(mid)] = center
+
+        return markers
+
+    def capture_state(self) -> Tuple[Optional[Tuple[int, int]], Dict[int, Tuple[float, float]], np.ndarray]:
+        """Capture frame and return (ball_pos, markers, frame)."""
+        if not self.cap:
+            return None, {}, np.zeros((480, 640, 3), dtype=np.uint8)
+
+        # Flush old frames
+        for _ in range(3):
+            self.cap.read()
+
+        ret, frame = self.cap.read()
+        if not ret:
+            return None, {}, np.zeros((480, 640, 3), dtype=np.uint8)
+
+        ball = self.detect_ball(frame)
+        markers = self.detect_markers(frame)
+
+        return ball, markers, frame
+
+    def distance(self, p1: Optional[Tuple[float, float]], p2: Optional[Tuple[float, float]]) -> float:
+        """Euclidean distance between two points."""
+        if p1 is None or p2 is None:
+            return float("inf")
+        return np.sqrt((p1[0] - p2[0])**2 + (p1[1] - p2[1])**2)
+
+    def twitch_test(
+        self,
+        servo_id: int,
+        neutral: int = 1500,
+        perturbation: int = 100,
+    ) -> Optional[TwitchResult]:
+        """
+        Perform a single twitch test on a servo.
+
+        1. Move to neutral
+        2. Measure distance(gripper, ball)
+        3. Move by +perturbation
+        4. Measure distance again
+        5. Determine if we got closer or further
+        """
+        print(f"\n  Testing servo {servo_id}...")
+
+        # Move to neutral first
+        self.set_servo(servo_id, neutral, duration=600)
+        time.sleep(0.5)
+
+        # Capture baseline
+        ball_before, markers_before, _ = self.capture_state()
+
+        if ball_before is None:
+            print(f"    Warning: Ball not detected")
+            return None
+
+        gripper_before = markers_before.get(self.gripper_marker_id)
+        if gripper_before is None:
+            # Try other arm markers as fallback
+            for marker_id in [15, 14, 13]:
+                if marker_id in markers_before:
+                    gripper_before = markers_before[marker_id]
+                    print(f"    Using marker M{marker_id} as reference")
+                    break
+
+        if gripper_before is None:
+            print(f"    Warning: No arm markers detected")
+            return None
+
+        dist_before = self.distance(gripper_before, ball_before)
+        print(f"    Before: gripper at {gripper_before[0]:.0f},{gripper_before[1]:.0f}, dist={dist_before:.1f}px")
+
+        # Perturbation: move servo by +perturbation
+        self.set_servo(servo_id, neutral + perturbation, duration=400)
+        time.sleep(0.4)
+
+        # Capture after perturbation
+        ball_after, markers_after, frame = self.capture_state()
+
+        if ball_after is None:
+            print(f"    Warning: Ball lost after move")
+            # Return to neutral
+            self.set_servo(servo_id, neutral, duration=400)
+            return None
+
+        # Find gripper marker (same ID as before if possible)
+        gripper_after = None
+        for marker_id in [15, 14, 13]:
+            if marker_id in markers_after:
+                gripper_after = markers_after[marker_id]
+                break
+
+        if gripper_after is None:
+            print(f"    Warning: Gripper marker lost after move (significant movement)")
+            # Return to neutral
+            self.set_servo(servo_id, neutral, duration=400)
+            # If marker was lost, the movement was significant - treat as "toward" if likely
+            return TwitchResult(
+                servo_id=servo_id,
+                initial_distance=dist_before,
+                final_distance=0,  # Unknown
+                delta=-dist_before,  # Assume moved toward (marker closer = can't see)
+                perturbation=perturbation,
+                direction="toward",  # Tentative
+                confidence=0.3,  # Low confidence since we can't verify
+            )
+
+        dist_after = self.distance(gripper_after, ball_after)
+        delta = dist_after - dist_before
+
+        print(f"    After:  gripper at {gripper_after[0]:.0f},{gripper_after[1]:.0f}, dist={dist_after:.1f}px")
+        print(f"    Delta:  {delta:+.1f}px")
+
+        # Return to neutral
+        self.set_servo(servo_id, neutral, duration=400)
+        time.sleep(0.3)
+
+        # Determine direction
+        # Negative delta = got closer = positive command moves toward target
+        if delta < -10:
+            direction = "toward"
+            confidence = min(1.0, abs(delta) / 50)
+        elif delta > 10:
+            direction = "away"
+            confidence = min(1.0, abs(delta) / 50)
+        else:
+            direction = "minimal"
+            confidence = 0.2
+
+        print(f"    Result: +{perturbation} PWM → {direction} (confidence: {confidence:.0%})")
+
+        return TwitchResult(
+            servo_id=servo_id,
+            initial_distance=dist_before,
+            final_distance=dist_after,
+            delta=delta,
+            perturbation=perturbation,
+            direction=direction,
+            confidence=confidence,
+        )
+
+    def run_calibration(self) -> Dict[int, DirectionMapping]:
+        """
+        Run direction calibration for all arm servos.
+
+        Returns dict of servo_id -> DirectionMapping
+        """
+        print("\n" + "=" * 60)
+        print("DIRECTION CALIBRATION (Twitch Test)")
+        print("=" * 60)
+        print("\nPurpose: Determine if positive servo values move")
+        print("         TOWARD or AWAY from the target.\n")
+
+        if not self.setup_camera():
+            return {}
+
+        if not self.setup_robot():
+            self.cleanup()
+            return {}
+
+        # Verify we can see ball and markers
+        print("Verifying visibility...")
+        ball, markers, frame = self.capture_state()
+
+        if ball is None:
+            print("\nERROR: Cannot detect ball.")
+            print("Please position a green ball in front of the robot.")
+            cv2.imwrite("/tmp/direction_cal_no_ball.jpg", frame)
+            print("Saved debug image to /tmp/direction_cal_no_ball.jpg")
+            self.cleanup()
+            return {}
+
+        print(f"  Ball detected at: {ball}")
+
+        gripper_marker = None
+        for m in [15, 14, 13]:
+            if m in markers:
+                gripper_marker = m
+                break
+
+        if gripper_marker is None:
+            print("\nERROR: No arm markers (M13, M14, M15) visible.")
+            print("Please position camera to see arm markers.")
+            cv2.imwrite("/tmp/direction_cal_no_markers.jpg", frame)
+            print("Saved debug image to /tmp/direction_cal_no_markers.jpg")
+            self.cleanup()
+            return {}
+
+        print(f"  Using marker M{gripper_marker} as gripper reference")
+        print(f"  Visible markers: {sorted(markers.keys())}")
+
+        # Save initial state image
+        cv2.circle(frame, ball, 20, (0, 255, 0), 3)
+        cv2.putText(frame, "BALL", (ball[0]+25, ball[1]),
+                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
+        gpos = markers[gripper_marker]
+        cv2.circle(frame, (int(gpos[0]), int(gpos[1])), 15, (255, 0, 255), 3)
+        cv2.putText(frame, f"M{gripper_marker}", (int(gpos[0])+20, int(gpos[1])),
+                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 255), 2)
+        cv2.imwrite("/tmp/direction_cal_initial.jpg", frame)
+
+        # Set neutral position for all arm servos
+        print("\nSetting neutral position...")
+        for servo_id in self.arm_servos:
+            self.set_servo(servo_id, 1500, duration=500)
+        time.sleep(1)
+
+        # Run twitch test on each servo
+        print("\nRunning twitch tests...")
+        direction_mappings: Dict[int, DirectionMapping] = {}
+
+        for servo_id in self.arm_servos:
+            result = self.twitch_test(servo_id)
+
+            if result:
+                self.results[servo_id] = result
+
+                # Convert to DirectionMapping
+                toward_direction = "positive" if result.direction == "toward" else "negative"
+                if result.direction == "minimal":
+                    toward_direction = "unknown"
+
+                mapping = DirectionMapping(
+                    servo_id=servo_id,
+                    toward_direction=toward_direction,
+                    confidence=result.confidence,
+                    delta_pixels=result.delta,
+                    calibrated_at=datetime.now().isoformat(),
+                )
+                direction_mappings[servo_id] = mapping
+
+        self.cleanup()
+
+        # Print summary
+        print("\n" + "=" * 60)
+        print("DIRECTION CALIBRATION RESULTS")
+        print("=" * 60)
+
+        for servo_id, mapping in direction_mappings.items():
+            result = self.results.get(servo_id)
+            delta_str = f"Δ={result.delta:+.1f}px" if result else ""
+            print(f"\n  Servo {servo_id}:")
+            print(f"    To move TOWARD target: use {mapping.toward_direction} values")
+            print(f"    Confidence: {mapping.confidence:.0%} {delta_str}")
+
+        # Save results
+        output_path = Path.home() / ".ate" / "direction_calibration" / f"{self.robot_name}.json"
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+
+        output = {
+            "robot_name": self.robot_name,
+            "calibrated_at": datetime.now().isoformat(),
+            "gripper_marker_id": self.gripper_marker_id,
+            "mappings": {
+                str(k): {
+                    "servo_id": v.servo_id,
+                    "toward_direction": v.toward_direction,
+                    "confidence": v.confidence,
+                    "delta_pixels": v.delta_pixels,
+                }
+                for k, v in direction_mappings.items()
+            },
+        }
+
+        with open(output_path, "w") as f:
+            json.dump(output, f, indent=2)
+
+        print(f"\nSaved to: {output_path}")
+
+        return direction_mappings
+
+    def cleanup(self) -> None:
+        """Release resources."""
+        if self.cap:
+            self.cap.release()
+        if self.ser:
+            self.ser.close()
+
+
+def run_direction_calibration(
+    serial_port: str,
+    robot_name: str = "robot",
+    gripper_marker_id: int = 15,
+    camera_index: int = 0,
+    arm_servos: Optional[List[int]] = None,
+) -> Dict[int, DirectionMapping]:
+    """
+    Run direction calibration and update calibration state.
+
+    This is the main entry point for the 'ate robot calibrate direction' command.
+    """
+    calibrator = DirectionCalibrator(
+        serial_port=serial_port,
+        robot_name=robot_name,
+        gripper_marker_id=gripper_marker_id,
+        camera_index=camera_index,
+        arm_servos=arm_servos,
+    )
+
+    mappings = calibrator.run_calibration()
+
+    if mappings:
+        # Update calibration state
+        state = CalibrationState.load(robot_name)
+        state.complete_direction_calibration(mappings)
+
+        print("\n" + "=" * 60)
+        print("CALIBRATION STATE UPDATED")
+        print("=" * 60)
+        state.print_status()
+
+    return mappings
+
+
+def load_direction_calibration(robot_name: str) -> Optional[Dict[int, DirectionMapping]]:
+    """Load saved direction calibration for a robot."""
+    path = Path.home() / ".ate" / "direction_calibration" / f"{robot_name}.json"
+
+    if not path.exists():
+        return None
+
+    try:
+        with open(path) as f:
+            data = json.load(f)
+
+        mappings = {}
+        for k, v in data.get("mappings", {}).items():
+            mappings[int(k)] = DirectionMapping(
+                servo_id=v["servo_id"],
+                toward_direction=v["toward_direction"],
+                confidence=v["confidence"],
+                delta_pixels=v["delta_pixels"],
+                calibrated_at=v.get("calibrated_at", ""),
+            )
+
+        return mappings
+
+    except Exception as e:
+        print(f"Error loading direction calibration: {e}")
+        return None
+
+
+def get_toward_direction(robot_name: str, servo_id: int) -> Optional[str]:
+    """
+    Get the direction that moves a servo toward the target.
+
+    Returns "positive" or "negative", or None if not calibrated.
+    """
+    mappings = load_direction_calibration(robot_name)
+    if not mappings or servo_id not in mappings:
+        return None
+
+    return mappings[servo_id].toward_direction