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

ate/detection/trash_detector.py

@@ -0,0 +1,322 @@
+"""
+Trash detection for cleanup robots.
+
+Combines multiple detection strategies:
+1. Color detection (bottles, cans, wrappers)
+2. Shape heuristics (round objects, rectangular objects)
+3. Location heuristics (ground-level objects)
+
+This is a practical starting point that works without ML models.
+For production, integrate with YOLO or similar.
+"""
+
+from dataclasses import dataclass
+from typing import List, Optional, Any, Dict, Set
+import time
+
+from .base import DetectorBase, Detection, BoundingBox
+from .color_detector import ColorDetector, ColorRange, PREDEFINED_COLORS, HAS_PIL
+
+
+# Trash-specific color ranges
+TRASH_COLORS: Dict[str, ColorRange] = {
+    # Plastic bottles (often blue/green/clear)
+    "plastic_blue": ColorRange("plastic", 90, 120, 30, 200, 100, 255),
+    "plastic_green": ColorRange("plastic", 40, 80, 30, 200, 100, 255),
+
+    # Metal cans (silver/gray with reflections)
+    "metal": ColorRange("metal", 0, 179, 0, 50, 100, 220),
+
+    # Wrappers/packaging (bright colors)
+    "wrapper_red": PREDEFINED_COLORS["red"],
+    "wrapper_yellow": PREDEFINED_COLORS["yellow"],
+    "wrapper_orange": PREDEFINED_COLORS["orange"],
+
+    # Paper/cardboard (brown/tan)
+    "cardboard": ColorRange("cardboard", 10, 25, 30, 150, 80, 200),
+}
+
+
+@dataclass
+class TrashItem:
+    """A detected piece of trash."""
+    detection: Detection
+    trash_type: str          # bottle, can, wrapper, paper, unknown
+    confidence: float        # Combined confidence
+    is_on_ground: bool       # Likely on ground level
+    size_category: str       # small, medium, large
+    priority: float          # Pickup priority (0-1)
+
+
+class TrashDetector(DetectorBase):
+    """
+    Specialized detector for trash/litter.
+
+    Designed for cleanup robots - prioritizes:
+    - Ground-level objects
+    - Common litter types (bottles, cans, wrappers)
+    - Pickable size objects
+
+    Example:
+        detector = TrashDetector()
+
+        # Detect all trash
+        items = detector.detect_trash(image)
+
+        # Find nearest pickable trash
+        nearest = detector.find_pickable(image)
+        if nearest:
+            print(f"Found {nearest.trash_type} at {nearest.detection.center}")
+    """
+
+    def __init__(
+        self,
+        min_area: int = 300,
+        max_area: int = 50000,
+        ground_threshold: float = 0.6,  # Bottom 60% of image
+        confidence_threshold: float = 0.3,
+    ):
+        """
+        Initialize trash detector.
+
+        Args:
+            min_area: Minimum object area in pixels
+            max_area: Maximum object area (too big = not trash)
+            ground_threshold: Fraction of image height considered "ground"
+            confidence_threshold: Minimum confidence to report
+        """
+        if not HAS_PIL:
+            raise ImportError(
+                "TrashDetector requires Pillow. Install with: pip install 'foodforthought-cli[detection]'"
+            )
+
+        self.min_area = min_area
+        self.max_area = max_area
+        self.ground_threshold = ground_threshold
+        self.confidence_threshold = confidence_threshold
+
+        # Color detector for actual detection
+        self._color_detector = ColorDetector(
+            min_area=min_area,
+            max_area=max_area,
+        )
+
+    def detect(self, image: Any, **kwargs) -> List[Detection]:
+        """Detect all colored objects (for base interface)."""
+        return self._color_detector.detect(image, **kwargs)
+
+    def detect_class(self, image: Any, class_name: str, **kwargs) -> List[Detection]:
+        """Detect objects of a specific class."""
+        if class_name == "trash":
+            items = self.detect_trash(image)
+            return [item.detection for item in items]
+        return self._color_detector.detect_class(image, class_name, **kwargs)
+
+    def detect_trash(self, image: Any) -> List[TrashItem]:
+        """
+        Detect trash items in image.
+
+        Returns list of TrashItem with classification and priority.
+        """
+        # Get image dimensions
+        width, height = self._get_dimensions(image)
+        if width == 0:
+            return []
+
+        # Detect with trash-specific colors
+        color_ranges = list(TRASH_COLORS.values())
+        detections = self._color_detector.detect(image, color_ranges=color_ranges)
+
+        # Convert to TrashItems with classification
+        items = []
+        for det in detections:
+            item = self._classify_trash(det, width, height)
+            if item.confidence >= self.confidence_threshold:
+                items.append(item)
+
+        # Sort by priority (highest first)
+        items.sort(key=lambda x: x.priority, reverse=True)
+
+        return items
+
+    def find_pickable(
+        self,
+        image: Any,
+        reference_point: tuple = None,
+        max_size: str = "large",
+    ) -> Optional[TrashItem]:
+        """
+        Find the best trash item to pick up.
+
+        Considers:
+        - Distance from reference point (or image center)
+        - Size (pickable range)
+        - Priority
+
+        Args:
+            image: Input image
+            reference_point: Reference point (default: bottom center)
+            max_size: Maximum size category to consider
+
+        Returns:
+            Best TrashItem to pick, or None
+        """
+        items = self.detect_trash(image)
+        if not items:
+            return None
+
+        # Filter by size
+        size_order = ["small", "medium", "large"]
+        max_idx = size_order.index(max_size) if max_size in size_order else 2
+        items = [i for i in items if size_order.index(i.size_category) <= max_idx]
+
+        if not items:
+            return None
+
+        # Default reference: bottom center of image
+        width, height = self._get_dimensions(image)
+        if reference_point is None:
+            reference_point = (width // 2, int(height * 0.9))
+
+        # Score each item
+        def score(item):
+            # Distance score (closer is better)
+            cx, cy = item.detection.center
+            dist = ((cx - reference_point[0])**2 + (cy - reference_point[1])**2)**0.5
+            max_dist = (width**2 + height**2)**0.5
+            dist_score = 1.0 - (dist / max_dist)
+
+            # Combine with priority
+            return 0.6 * item.priority + 0.4 * dist_score
+
+        return max(items, key=score)
+
+    def _classify_trash(self, det: Detection, width: int, height: int) -> TrashItem:
+        """Classify a detection as a trash type."""
+        bbox = det.bbox
+        cx, cy = bbox.center
+        area = bbox.area
+
+        # Determine if on ground (in lower portion of image)
+        is_on_ground = cy > height * (1 - self.ground_threshold)
+
+        # Classify size
+        image_area = width * height
+        area_ratio = area / image_area
+        if area_ratio < 0.01:
+            size_category = "small"
+        elif area_ratio < 0.05:
+            size_category = "medium"
+        else:
+            size_category = "large"
+
+        # Classify trash type based on color and shape
+        color = det.metadata.get("color", det.label)
+        aspect = bbox.width / max(bbox.height, 1)
+
+        if "plastic" in color or det.label in ["blue", "green"]:
+            if 0.3 < aspect < 0.6:  # Tall and thin = bottle
+                trash_type = "bottle"
+            else:
+                trash_type = "plastic"
+        elif "metal" in color or det.label == "gray":
+            if 0.7 < aspect < 1.3:  # Round-ish = can
+                trash_type = "can"
+            else:
+                trash_type = "metal"
+        elif "cardboard" in color:
+            trash_type = "paper"
+        elif det.label in ["red", "yellow", "orange"]:
+            trash_type = "wrapper"
+        else:
+            trash_type = "unknown"
+
+        # Calculate priority
+        # Higher priority for:
+        # - Objects on ground
+        # - Pickable size
+        # - Higher confidence
+        priority = det.confidence
+
+        if is_on_ground:
+            priority *= 1.3
+
+        if size_category == "medium":
+            priority *= 1.2
+        elif size_category == "small":
+            priority *= 1.1
+        # Large items get no bonus (harder to pick)
+
+        # Certain types are easier to pick
+        if trash_type in ["bottle", "can"]:
+            priority *= 1.1
+
+        priority = min(priority, 1.0)
+
+        return TrashItem(
+            detection=det,
+            trash_type=trash_type,
+            confidence=det.confidence,
+            is_on_ground=is_on_ground,
+            size_category=size_category,
+            priority=priority,
+        )
+
+    def _get_dimensions(self, image: Any) -> tuple:
+        """Get image dimensions."""
+        if hasattr(image, 'width') and hasattr(image, 'height'):
+            return (image.width, image.height)
+
+        try:
+            from PIL import Image as PILImage
+            if isinstance(image, PILImage.Image):
+                return image.size
+        except ImportError:
+            pass
+
+        return (640, 480)  # Default VGA
+
+
+def demo_detection(camera_or_image: Any = None):
+    """
+    Demo the trash detector.
+
+    Can be run standalone for testing:
+        python -m ate.detection.trash_detector
+    """
+    detector = TrashDetector()
+
+    if camera_or_image is None:
+        print("TrashDetector Demo")
+        print("=" * 40)
+        print("\nUsage:")
+        print("  from ate.detection import TrashDetector")
+        print("  detector = TrashDetector()")
+        print("  items = detector.detect_trash(image)")
+        print("  pickable = detector.find_pickable(image)")
+        print("\nTrash types detected:")
+        print("  - bottle (plastic bottles)")
+        print("  - can (metal cans)")
+        print("  - wrapper (colorful packaging)")
+        print("  - paper (cardboard, paper)")
+        print("  - plastic (other plastic items)")
+        print("  - metal (other metal items)")
+        print("  - unknown (unclassified)")
+        return
+
+    # Detect from provided image
+    items = detector.detect_trash(camera_or_image)
+
+    print(f"\nFound {len(items)} trash items:")
+    for item in items:
+        cx, cy = item.detection.center
+        print(f"  {item.trash_type}: ({cx}, {cy}) - {item.size_category}, priority={item.priority:.2f}")
+
+    pickable = detector.find_pickable(camera_or_image)
+    if pickable:
+        cx, cy = pickable.detection.center
+        print(f"\nBest to pick: {pickable.trash_type} at ({cx}, {cy})")
+
+
+if __name__ == "__main__":
+    demo_detection()

ate/drivers/__init__.py

@@ -4,24 +4,36 @@ Hardware drivers for specific robots.
 Each driver implements the appropriate interfaces from ate.interfaces.
 
 Example:
-    from ate.drivers import MechDogDriver
+    from ate.drivers import MechDogDriver, MechDogConfig
 
-    # Connect to MechDog
-    dog = MechDogDriver(port="/dev/cu.usbserial-10")
+    # Connect to MechDog with camera
+    config = MechDogConfig(
+        port="/dev/cu.usbserial-10",
+        has_camera=True,
+        camera_ip="192.168.1.100"
+    )
+    dog = MechDogDriver(config=config)
     dog.connect()
 
     # Use through abstract interface
     dog.stand()
     dog.walk(Vector3.forward(), speed=0.3)
     dog.set_body_height(0.15)
-    dog.stop()
 
-    # Disconnect
+    # Capture image from visual module
+    image = dog.get_image()
+
+    dog.stop()
     dog.disconnect()
 """
 
-from .mechdog import MechDogDriver
+from .mechdog import MechDogDriver, MechDogConfig
+from .wifi_camera import WiFiCamera, WiFiCameraConfig, discover_cameras
 
 __all__ = [
     "MechDogDriver",
+    "MechDogConfig",
+    "WiFiCamera",
+    "WiFiCameraConfig",
+    "discover_cameras",
 ]