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

ate/detection/color_detector.py

@@ -0,0 +1,399 @@
+"""
+Color-based object detection.
+
+Simple but effective detection using color segmentation.
+Works without ML dependencies - just needs PIL/Pillow.
+
+Great for:
+- Detecting colored objects (red balls, blue bins, etc.)
+- Prototyping before using ML models
+- Real-time detection on limited hardware
+"""
+
+from dataclasses import dataclass
+from typing import List, Tuple, Optional, Any, Dict
+import io
+
+from .base import DetectorBase, Detection, BoundingBox
+
+try:
+    from PIL import Image as PILImage
+    HAS_PIL = True
+except ImportError:
+    HAS_PIL = False
+
+
+@dataclass
+class ColorRange:
+    """
+    HSV color range for detection.
+
+    HSV is better than RGB for color detection because
+    it separates color (hue) from brightness.
+    """
+    name: str
+    h_min: int  # Hue min (0-179)
+    h_max: int  # Hue max (0-179)
+    s_min: int  # Saturation min (0-255)
+    s_max: int  # Saturation max (0-255)
+    v_min: int  # Value min (0-255)
+    v_max: int  # Value max (0-255)
+
+
+# Predefined color ranges (HSV)
+PREDEFINED_COLORS: Dict[str, ColorRange] = {
+    "red": ColorRange("red", 0, 10, 100, 255, 100, 255),
+    "red2": ColorRange("red2", 160, 179, 100, 255, 100, 255),  # Red wraps around
+    "orange": ColorRange("orange", 11, 25, 100, 255, 100, 255),
+    "yellow": ColorRange("yellow", 26, 35, 100, 255, 100, 255),
+    "green": ColorRange("green", 36, 85, 50, 255, 50, 255),
+    "blue": ColorRange("blue", 86, 125, 50, 255, 50, 255),
+    "purple": ColorRange("purple", 126, 155, 50, 255, 50, 255),
+    "pink": ColorRange("pink", 156, 165, 50, 255, 100, 255),
+    "white": ColorRange("white", 0, 179, 0, 30, 200, 255),
+    "black": ColorRange("black", 0, 179, 0, 255, 0, 50),
+    "gray": ColorRange("gray", 0, 179, 0, 30, 50, 200),
+}
+
+
+class ColorDetector(DetectorBase):
+    """
+    Detect objects by color using HSV thresholding.
+
+    Example:
+        detector = ColorDetector()
+
+        # Detect red objects
+        detections = detector.detect(image, target_colors=["red"])
+
+        # Or use custom color range
+        custom = ColorRange("bright_green", 40, 80, 150, 255, 150, 255)
+        detections = detector.detect(image, color_ranges=[custom])
+    """
+
+    def __init__(
+        self,
+        min_area: int = 500,
+        max_area: int = 100000,
+        merge_distance: int = 20,
+    ):
+        """
+        Initialize color detector.
+
+        Args:
+            min_area: Minimum detection area in pixels
+            max_area: Maximum detection area in pixels
+            merge_distance: Distance to merge nearby detections
+        """
+        if not HAS_PIL:
+            raise ImportError("ColorDetector requires Pillow. Install with: pip install Pillow")
+
+        self.min_area = min_area
+        self.max_area = max_area
+        self.merge_distance = merge_distance
+
+    def detect(
+        self,
+        image: Any,
+        target_colors: List[str] = None,
+        color_ranges: List[ColorRange] = None,
+        **kwargs
+    ) -> List[Detection]:
+        """
+        Detect colored objects.
+
+        Args:
+            image: Input image (Image dataclass, PIL Image, or bytes)
+            target_colors: List of color names ["red", "blue", ...]
+            color_ranges: Custom ColorRange objects
+
+        Returns:
+            List of Detection objects
+        """
+        # Convert image to PIL
+        pil_image = self._to_pil(image)
+        if pil_image is None:
+            return []
+
+        # Build list of color ranges to detect
+        ranges = []
+        if target_colors:
+            for color in target_colors:
+                if color in PREDEFINED_COLORS:
+                    ranges.append(PREDEFINED_COLORS[color])
+                    # Handle red wraparound
+                    if color == "red":
+                        ranges.append(PREDEFINED_COLORS["red2"])
+
+        if color_ranges:
+            ranges.extend(color_ranges)
+
+        if not ranges:
+            # Default: detect common trash colors
+            ranges = [
+                PREDEFINED_COLORS["red"],
+                PREDEFINED_COLORS["red2"],
+                PREDEFINED_COLORS["blue"],
+                PREDEFINED_COLORS["green"],
+                PREDEFINED_COLORS["yellow"],
+            ]
+
+        # Detect for each color range
+        detections = []
+        for color_range in ranges:
+            dets = self._detect_color(pil_image, color_range)
+            detections.extend(dets)
+
+        # Merge overlapping detections
+        detections = self._merge_detections(detections)
+
+        return detections
+
+    def detect_class(self, image: Any, class_name: str, **kwargs) -> List[Detection]:
+        """Detect objects of a specific color class."""
+        return self.detect(image, target_colors=[class_name], **kwargs)
+
+    def _to_pil(self, image: Any) -> Optional[PILImage.Image]:
+        """Convert various image formats to PIL Image."""
+        if isinstance(image, PILImage.Image):
+            return image
+
+        # Handle our Image dataclass
+        if hasattr(image, 'data') and hasattr(image, 'width'):
+            if image.width == 0 or image.height == 0:
+                return None
+
+            if image.encoding == "jpeg":
+                return PILImage.open(io.BytesIO(image.data))
+            elif image.encoding == "rgb8":
+                return PILImage.frombytes("RGB", (image.width, image.height), image.data)
+            else:
+                # Try to decode as JPEG
+                try:
+                    return PILImage.open(io.BytesIO(image.data))
+                except Exception:
+                    return None
+
+        # Handle raw bytes (assume JPEG)
+        if isinstance(image, bytes):
+            try:
+                return PILImage.open(io.BytesIO(image))
+            except Exception:
+                return None
+
+        return None
+
+    def _detect_color(self, image: PILImage.Image, color_range: ColorRange) -> List[Detection]:
+        """
+        Detect regions matching a color range.
+
+        Uses a simple but effective approach:
+        1. Convert to HSV
+        2. Threshold by color range
+        3. Find connected components
+        4. Filter by size
+        """
+        # Convert to RGB if needed
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+
+        width, height = image.size
+        pixels = image.load()
+
+        # Create binary mask
+        mask = [[False] * width for _ in range(height)]
+
+        for y in range(height):
+            for x in range(width):
+                r, g, b = pixels[x, y]
+                h, s, v = self._rgb_to_hsv(r, g, b)
+
+                # Check if in range
+                in_range = (
+                    color_range.s_min <= s <= color_range.s_max and
+                    color_range.v_min <= v <= color_range.v_max
+                )
+
+                # Handle hue wraparound for red
+                if color_range.h_min <= color_range.h_max:
+                    in_range = in_range and color_range.h_min <= h <= color_range.h_max
+                else:
+                    in_range = in_range and (h >= color_range.h_min or h <= color_range.h_max)
+
+                mask[y][x] = in_range
+
+        # Find connected components (simple flood fill)
+        visited = [[False] * width for _ in range(height)]
+        components = []
+
+        for y in range(height):
+            for x in range(width):
+                if mask[y][x] and not visited[y][x]:
+                    # Flood fill to find component
+                    component = self._flood_fill(mask, visited, x, y, width, height)
+                    if component:
+                        components.append(component)
+
+        # Convert components to detections
+        detections = []
+        for points in components:
+            area = len(points)
+            if area < self.min_area or area > self.max_area:
+                continue
+
+            # Calculate bounding box
+            min_x = min(p[0] for p in points)
+            max_x = max(p[0] for p in points)
+            min_y = min(p[1] for p in points)
+            max_y = max(p[1] for p in points)
+
+            bbox = BoundingBox(
+                x=min_x,
+                y=min_y,
+                width=max_x - min_x + 1,
+                height=max_y - min_y + 1,
+            )
+
+            # Confidence based on area and aspect ratio
+            aspect = bbox.width / max(bbox.height, 1)
+            aspect_score = 1.0 - min(abs(aspect - 1.0), 1.0)  # Prefer square-ish
+            area_score = min(area / self.max_area, 1.0)
+            confidence = 0.5 + 0.3 * aspect_score + 0.2 * area_score
+
+            detections.append(Detection(
+                label=color_range.name.replace("2", ""),  # Remove "red2" -> "red"
+                confidence=confidence,
+                bbox=bbox,
+                metadata={"area": area, "color": color_range.name},
+            ))
+
+        return detections
+
+    def _flood_fill(
+        self,
+        mask: List[List[bool]],
+        visited: List[List[bool]],
+        start_x: int,
+        start_y: int,
+        width: int,
+        height: int,
+        max_points: int = 50000,
+    ) -> List[Tuple[int, int]]:
+        """Flood fill to find connected component."""
+        stack = [(start_x, start_y)]
+        points = []
+
+        while stack and len(points) < max_points:
+            x, y = stack.pop()
+
+            if x < 0 or x >= width or y < 0 or y >= height:
+                continue
+            if visited[y][x] or not mask[y][x]:
+                continue
+
+            visited[y][x] = True
+            points.append((x, y))
+
+            # Add neighbors (4-connected)
+            stack.append((x + 1, y))
+            stack.append((x - 1, y))
+            stack.append((x, y + 1))
+            stack.append((x, y - 1))
+
+        return points
+
+    def _merge_detections(self, detections: List[Detection]) -> List[Detection]:
+        """Merge overlapping detections of the same color."""
+        if len(detections) <= 1:
+            return detections
+
+        merged = []
+        used = [False] * len(detections)
+
+        for i, det1 in enumerate(detections):
+            if used[i]:
+                continue
+
+            # Find all overlapping detections of same label
+            group = [det1]
+            used[i] = True
+
+            for j, det2 in enumerate(detections[i + 1:], i + 1):
+                if used[j]:
+                    continue
+                if det1.label != det2.label:
+                    continue
+
+                # Check if close enough to merge
+                dist = self._box_distance(det1.bbox, det2.bbox)
+                if dist < self.merge_distance:
+                    group.append(det2)
+                    used[j] = True
+
+            # Merge group into single detection
+            if len(group) == 1:
+                merged.append(det1)
+            else:
+                merged.append(self._merge_group(group))
+
+        return merged
+
+    def _box_distance(self, b1: BoundingBox, b2: BoundingBox) -> float:
+        """Calculate distance between bounding boxes."""
+        c1 = b1.center
+        c2 = b2.center
+        return ((c1[0] - c2[0]) ** 2 + (c1[1] - c2[1]) ** 2) ** 0.5
+
+    def _merge_group(self, group: List[Detection]) -> Detection:
+        """Merge multiple detections into one."""
+        # Calculate combined bounding box
+        min_x = min(d.bbox.x for d in group)
+        min_y = min(d.bbox.y for d in group)
+        max_x = max(d.bbox.x + d.bbox.width for d in group)
+        max_y = max(d.bbox.y + d.bbox.height for d in group)
+
+        bbox = BoundingBox(
+            x=min_x,
+            y=min_y,
+            width=max_x - min_x,
+            height=max_y - min_y,
+        )
+
+        # Average confidence
+        confidence = sum(d.confidence for d in group) / len(group)
+
+        # Sum area
+        total_area = sum(d.metadata.get("area", 0) for d in group)
+
+        return Detection(
+            label=group[0].label,
+            confidence=confidence,
+            bbox=bbox,
+            metadata={"area": total_area, "merged_count": len(group)},
+        )
+
+    def _rgb_to_hsv(self, r: int, g: int, b: int) -> Tuple[int, int, int]:
+        """Convert RGB (0-255) to HSV (H: 0-179, S: 0-255, V: 0-255)."""
+        r, g, b = r / 255.0, g / 255.0, b / 255.0
+        mx = max(r, g, b)
+        mn = min(r, g, b)
+        diff = mx - mn
+
+        # Hue
+        if diff == 0:
+            h = 0
+        elif mx == r:
+            h = (60 * ((g - b) / diff) + 360) % 360
+        elif mx == g:
+            h = (60 * ((b - r) / diff) + 120) % 360
+        else:
+            h = (60 * ((r - g) / diff) + 240) % 360
+
+        # Saturation
+        s = 0 if mx == 0 else (diff / mx)
+
+        # Value
+        v = mx
+
+        # Convert to OpenCV-style ranges
+        return int(h / 2), int(s * 255), int(v * 255)