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

ate/robot/target_calibration.py

@@ -0,0 +1,500 @@
+"""
+Target Detection Calibration
+
+Interactive tool to tune color detection for specific targets,
+with robot self-masking to ignore the robot's own features.
+
+The Problem:
+-----------
+Generic ball detection picked up:
+- Blue LEDs on robot
+- Cyan/teal reflections
+- ArUco marker borders
+
+This caused erratic position readings, breaking visual servoing.
+
+The Solution:
+------------
+1. Interactive HSV tuning with live preview
+2. Robot self-mask using ArUco marker positions
+3. Target profiles saved for reuse
+4. Multiple target types supported
+
+Usage:
+    from ate.robot.target_calibration import TargetCalibrator
+
+    cal = TargetCalibrator(camera_index=0)
+    profile = cal.calibrate_target("green_ball")
+    profile.save("~/.ate/targets/green_ball.json")
+
+    # Later use
+    profile = TargetProfile.load("green_ball")
+    position = profile.detect(frame)
+"""
+
+import cv2
+import numpy as np
+import json
+from dataclasses import dataclass, field, asdict
+from datetime import datetime
+from pathlib import Path
+from typing import Optional, Dict, List, Tuple, Callable
+import time
+
+
+@dataclass
+class HSVRange:
+    """HSV color range for detection."""
+    h_low: int = 0
+    h_high: int = 180
+    s_low: int = 0
+    s_high: int = 255
+    v_low: int = 0
+    v_high: int = 255
+
+    def to_arrays(self) -> Tuple[np.ndarray, np.ndarray]:
+        """Convert to OpenCV format arrays."""
+        lower = np.array([self.h_low, self.s_low, self.v_low])
+        upper = np.array([self.h_high, self.s_high, self.v_high])
+        return lower, upper
+
+    @classmethod
+    def from_dict(cls, d: dict) -> "HSVRange":
+        return cls(**d)
+
+
+@dataclass
+class TargetProfile:
+    """
+    Calibrated detection profile for a target type.
+
+    Includes HSV ranges, size constraints, and optional mask regions.
+    """
+    name: str
+    hsv_range: HSVRange
+    min_area: int = 500
+    max_area: int = 50000
+    min_circularity: float = 0.3  # 0-1, 1 = perfect circle
+    mask_regions: List[Tuple[int, int, int, int]] = field(default_factory=list)  # [(x, y, w, h), ...]
+    robot_mask_enabled: bool = True
+    created_at: str = field(default_factory=lambda: datetime.now().isoformat())
+    samples_collected: int = 0
+
+    def detect(
+        self,
+        frame: np.ndarray,
+        robot_markers: Optional[Dict[int, Tuple[float, float]]] = None,
+    ) -> Optional[Tuple[int, int]]:
+        """
+        Detect target in frame, return (x, y) or None.
+
+        Args:
+            frame: BGR image
+            robot_markers: Optional dict of ArUco marker positions to mask
+        """
+        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
+
+        # Create mask from HSV range
+        lower, upper = self.hsv_range.to_arrays()
+        mask = cv2.inRange(hsv, lower, upper)
+
+        # Apply static mask regions
+        for (x, y, w, h) in self.mask_regions:
+            mask[y:y+h, x:x+w] = 0
+
+        # Apply robot self-mask using ArUco markers
+        if self.robot_mask_enabled and robot_markers:
+            mask = self._apply_robot_mask(mask, robot_markers)
+
+        # Find contours
+        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+
+        if not contours:
+            return None
+
+        # Filter by area and circularity
+        valid_contours = []
+        for contour in contours:
+            area = cv2.contourArea(contour)
+            if area < self.min_area or area > self.max_area:
+                continue
+
+            # Calculate circularity
+            perimeter = cv2.arcLength(contour, True)
+            if perimeter > 0:
+                circularity = 4 * np.pi * area / (perimeter * perimeter)
+                if circularity < self.min_circularity:
+                    continue
+
+            valid_contours.append((contour, area))
+
+        if not valid_contours:
+            return None
+
+        # Return center of largest valid contour
+        largest = max(valid_contours, key=lambda x: x[1])[0]
+        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 _apply_robot_mask(
+        self,
+        mask: np.ndarray,
+        markers: Dict[int, Tuple[float, float]],
+        radius: int = 80,
+    ) -> np.ndarray:
+        """Mask out regions around detected ArUco markers (robot parts)."""
+        for marker_id, (mx, my) in markers.items():
+            # Mask a circular region around each marker
+            cv2.circle(mask, (int(mx), int(my)), radius, 0, -1)
+        return mask
+
+    def save(self, path: Optional[str] = None) -> str:
+        """Save profile to JSON."""
+        if path is None:
+            save_dir = Path.home() / ".ate" / "targets"
+            save_dir.mkdir(parents=True, exist_ok=True)
+            path = str(save_dir / f"{self.name}.json")
+
+        data = {
+            "name": self.name,
+            "hsv_range": asdict(self.hsv_range),
+            "min_area": self.min_area,
+            "max_area": self.max_area,
+            "min_circularity": self.min_circularity,
+            "mask_regions": self.mask_regions,
+            "robot_mask_enabled": self.robot_mask_enabled,
+            "created_at": self.created_at,
+            "samples_collected": self.samples_collected,
+        }
+
+        with open(path, "w") as f:
+            json.dump(data, f, indent=2)
+
+        return path
+
+    @classmethod
+    def load(cls, name_or_path: str) -> Optional["TargetProfile"]:
+        """Load profile from JSON."""
+        if "/" in name_or_path or name_or_path.endswith(".json"):
+            path = Path(name_or_path)
+        else:
+            path = Path.home() / ".ate" / "targets" / f"{name_or_path}.json"
+
+        if not path.exists():
+            return None
+
+        with open(path) as f:
+            data = json.load(f)
+
+        return cls(
+            name=data["name"],
+            hsv_range=HSVRange.from_dict(data["hsv_range"]),
+            min_area=data.get("min_area", 500),
+            max_area=data.get("max_area", 50000),
+            min_circularity=data.get("min_circularity", 0.3),
+            mask_regions=data.get("mask_regions", []),
+            robot_mask_enabled=data.get("robot_mask_enabled", True),
+            created_at=data.get("created_at", ""),
+            samples_collected=data.get("samples_collected", 0),
+        )
+
+
+class TargetCalibrator:
+    """
+    Interactive target calibration tool.
+
+    Provides live preview of detection with HSV tuning sliders.
+    """
+
+    def __init__(
+        self,
+        camera_index: int = 0,
+        aruco_dict: int = cv2.aruco.DICT_4X4_50,
+    ):
+        self.camera_index = camera_index
+        self.cap: Optional[cv2.VideoCapture] = None
+
+        # ArUco setup for robot masking
+        aruco_dictionary = cv2.aruco.getPredefinedDictionary(aruco_dict)
+        aruco_params = cv2.aruco.DetectorParameters()
+        self.aruco_detector = cv2.aruco.ArucoDetector(aruco_dictionary, aruco_params)
+
+        # Current HSV values (will be set by sliders)
+        self.hsv_range = HSVRange(
+            h_low=70, h_high=110,  # Blue-green default
+            s_low=50, s_high=255,
+            v_low=50, v_high=255,
+        )
+
+        # Samples for averaging
+        self.hsv_samples: List[Tuple[int, int, int]] = []
+
+    def setup_camera(self) -> bool:
+        """Initialize webcam."""
+        self.cap = cv2.VideoCapture(self.camera_index)
+        if not self.cap.isOpened():
+            return False
+
+        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
+        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
+
+        # Warm up
+        for _ in range(10):
+            self.cap.read()
+
+        return True
+
+    def detect_markers(self, frame: np.ndarray) -> Dict[int, Tuple[float, float]]:
+        """Detect ArUco markers for robot masking."""
+        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 sample_hsv_at_click(self, frame: np.ndarray, x: int, y: int, radius: int = 10):
+        """Sample HSV values from a region around click point."""
+        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
+
+        # Get region
+        y1 = max(0, y - radius)
+        y2 = min(frame.shape[0], y + radius)
+        x1 = max(0, x - radius)
+        x2 = min(frame.shape[1], x + radius)
+
+        region = hsv[y1:y2, x1:x2]
+
+        # Calculate median HSV values
+        h_med = int(np.median(region[:, :, 0]))
+        s_med = int(np.median(region[:, :, 1]))
+        v_med = int(np.median(region[:, :, 2]))
+
+        self.hsv_samples.append((h_med, s_med, v_med))
+
+        return (h_med, s_med, v_med)
+
+    def compute_hsv_range_from_samples(self, margin: int = 15) -> HSVRange:
+        """Compute HSV range from collected samples."""
+        if not self.hsv_samples:
+            return self.hsv_range
+
+        h_vals = [s[0] for s in self.hsv_samples]
+        s_vals = [s[1] for s in self.hsv_samples]
+        v_vals = [s[2] for s in self.hsv_samples]
+
+        return HSVRange(
+            h_low=max(0, min(h_vals) - margin),
+            h_high=min(180, max(h_vals) + margin),
+            s_low=max(0, min(s_vals) - margin),
+            s_high=min(255, max(s_vals) + margin),
+            v_low=max(0, min(v_vals) - margin),
+            v_high=min(255, max(v_vals) + margin),
+        )
+
+    def calibrate_interactive(self, target_name: str) -> Optional[TargetProfile]:
+        """
+        Run interactive calibration with OpenCV windows.
+
+        Instructions:
+        - Click on target object to sample HSV values
+        - Adjust sliders if needed
+        - Press 's' to save and exit
+        - Press 'c' to clear samples
+        - Press 'q' to quit without saving
+        """
+        if not self.setup_camera():
+            print("ERROR: Could not open camera")
+            return None
+
+        window_name = f"Target Calibration: {target_name}"
+        mask_window = "Detection Mask"
+
+        cv2.namedWindow(window_name)
+        cv2.namedWindow(mask_window)
+
+        # Mouse callback for sampling
+        current_frame = [None]
+
+        def mouse_callback(event, x, y, flags, param):
+            if event == cv2.EVENT_LBUTTONDOWN and current_frame[0] is not None:
+                hsv = self.sample_hsv_at_click(current_frame[0], x, y)
+                print(f"Sampled HSV at ({x}, {y}): H={hsv[0]}, S={hsv[1]}, V={hsv[2]}")
+                # Update range from samples
+                self.hsv_range = self.compute_hsv_range_from_samples()
+
+        cv2.setMouseCallback(window_name, mouse_callback)
+
+        # Create trackbars
+        def nothing(x):
+            pass
+
+        cv2.createTrackbar("H Low", window_name, self.hsv_range.h_low, 180, nothing)
+        cv2.createTrackbar("H High", window_name, self.hsv_range.h_high, 180, nothing)
+        cv2.createTrackbar("S Low", window_name, self.hsv_range.s_low, 255, nothing)
+        cv2.createTrackbar("S High", window_name, self.hsv_range.s_high, 255, nothing)
+        cv2.createTrackbar("V Low", window_name, self.hsv_range.v_low, 255, nothing)
+        cv2.createTrackbar("V High", window_name, self.hsv_range.v_high, 255, nothing)
+        cv2.createTrackbar("Robot Mask", window_name, 1, 1, nothing)
+
+        print(f"\nTarget Calibration: {target_name}")
+        print("=" * 50)
+        print("Click on the target object to sample HSV values")
+        print("Adjust sliders to fine-tune detection")
+        print("Press 's' to save, 'c' to clear samples, 'q' to quit")
+        print()
+
+        robot_mask_enabled = True
+
+        while True:
+            ret, frame = self.cap.read()
+            if not ret:
+                continue
+
+            current_frame[0] = frame.copy()
+
+            # Read trackbar values
+            self.hsv_range.h_low = cv2.getTrackbarPos("H Low", window_name)
+            self.hsv_range.h_high = cv2.getTrackbarPos("H High", window_name)
+            self.hsv_range.s_low = cv2.getTrackbarPos("S Low", window_name)
+            self.hsv_range.s_high = cv2.getTrackbarPos("S High", window_name)
+            self.hsv_range.v_low = cv2.getTrackbarPos("V Low", window_name)
+            self.hsv_range.v_high = cv2.getTrackbarPos("V High", window_name)
+            robot_mask_enabled = cv2.getTrackbarPos("Robot Mask", window_name) == 1
+
+            # Detect markers for masking
+            markers = self.detect_markers(frame)
+
+            # Create profile for detection
+            profile = TargetProfile(
+                name=target_name,
+                hsv_range=self.hsv_range,
+                robot_mask_enabled=robot_mask_enabled,
+                samples_collected=len(self.hsv_samples),
+            )
+
+            # Detect target
+            target_pos = profile.detect(frame, markers if robot_mask_enabled else None)
+
+            # Create visualization
+            display = frame.copy()
+
+            # Draw marker masks
+            if robot_mask_enabled:
+                for mid, (mx, my) in markers.items():
+                    cv2.circle(display, (int(mx), int(my)), 80, (128, 128, 128), 2)
+                    cv2.putText(display, f"MASK M{mid}", (int(mx)-30, int(my)-85),
+                               cv2.FONT_HERSHEY_SIMPLEX, 0.5, (128, 128, 128), 1)
+
+            # Draw detected target
+            if target_pos:
+                cv2.circle(display, target_pos, 30, (0, 255, 0), 3)
+                cv2.putText(display, f"TARGET ({target_pos[0]}, {target_pos[1]})",
+                           (target_pos[0] + 35, target_pos[1]),
+                           cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
+
+            # Draw HSV info
+            info = f"HSV: H[{self.hsv_range.h_low}-{self.hsv_range.h_high}] "
+            info += f"S[{self.hsv_range.s_low}-{self.hsv_range.s_high}] "
+            info += f"V[{self.hsv_range.v_low}-{self.hsv_range.v_high}]"
+            cv2.putText(display, info, (10, 30),
+                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
+
+            cv2.putText(display, f"Samples: {len(self.hsv_samples)}", (10, 60),
+                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
+
+            # Show detection mask
+            hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
+            lower, upper = self.hsv_range.to_arrays()
+            mask = cv2.inRange(hsv, lower, upper)
+
+            if robot_mask_enabled:
+                for mid, (mx, my) in markers.items():
+                    cv2.circle(mask, (int(mx), int(my)), 80, 0, -1)
+
+            cv2.imshow(window_name, display)
+            cv2.imshow(mask_window, mask)
+
+            key = cv2.waitKey(1) & 0xFF
+
+            if key == ord('s'):
+                # Save and exit
+                profile.samples_collected = len(self.hsv_samples)
+                path = profile.save()
+                print(f"\nSaved target profile to: {path}")
+                self.cap.release()
+                cv2.destroyAllWindows()
+                return profile
+
+            elif key == ord('c'):
+                # Clear samples
+                self.hsv_samples = []
+                print("Cleared HSV samples")
+
+            elif key == ord('q'):
+                # Quit without saving
+                print("\nCancelled - profile not saved")
+                self.cap.release()
+                cv2.destroyAllWindows()
+                return None
+
+        self.cap.release()
+        cv2.destroyAllWindows()
+        return None
+
+    def calibrate_non_interactive(
+        self,
+        target_name: str,
+        hsv_range: HSVRange,
+        robot_mask_enabled: bool = True,
+    ) -> TargetProfile:
+        """Create target profile without interactive calibration."""
+        return TargetProfile(
+            name=target_name,
+            hsv_range=hsv_range,
+            robot_mask_enabled=robot_mask_enabled,
+        )
+
+
+def run_target_calibration(target_name: str, camera_index: int = 0) -> Optional[TargetProfile]:
+    """
+    Run interactive target calibration.
+
+    Entry point for CLI command.
+    """
+    calibrator = TargetCalibrator(camera_index=camera_index)
+    return calibrator.calibrate_interactive(target_name)
+
+
+def list_target_profiles() -> List[str]:
+    """List saved target profiles."""
+    target_dir = Path.home() / ".ate" / "targets"
+    if not target_dir.exists():
+        return []
+    return [p.stem for p in target_dir.glob("*.json")]
+
+
+def detect_with_profile(
+    frame: np.ndarray,
+    profile_name: str,
+    robot_markers: Optional[Dict[int, Tuple[float, float]]] = None,
+) -> Optional[Tuple[int, int]]:
+    """
+    Detect target using saved profile.
+
+    Convenience function for use in behaviors.
+    """
+    profile = TargetProfile.load(profile_name)
+    if not profile:
+        return None
+    return profile.detect(frame, robot_markers)