eye-cv 1.0.0__py3-none-any.whl

eye/detection/detections_enhanced.py

@@ -0,0 +1,392 @@
+"""Core detection management classes with universal auto-conversion."""
+
+import numpy as np
+from typing import Optional, Dict, List, Tuple, Any, Union
+from dataclasses import dataclass, field
+
+
+@dataclass
+class Detections:
+    """Container for object detections with optimized operations.
+    
+    Design notes:
+    - Immutable operations (returns new instance instead of modifying)
+    - Built-in caching for expensive computations
+    - Native support for segmentation masks
+    - Better indexing and slicing
+    """
+    
+    xyxy: np.ndarray  # Shape: (N, 4) - [x1, y1, x2, y2]
+    confidence: Optional[np.ndarray] = None  # Shape: (N,)
+    class_id: Optional[np.ndarray] = None  # Shape: (N,)
+    tracker_id: Optional[np.ndarray] = None  # Shape: (N,)
+    mask: Optional[np.ndarray] = None  # Shape: (N, H, W) - segmentation masks
+    data: Dict[str, Any] = field(default_factory=dict)  # Additional metadata
+    
+    # Cache for expensive computations
+    _area_cache: Optional[np.ndarray] = field(default=None, repr=False, compare=False)
+    _center_cache: Optional[np.ndarray] = field(default=None, repr=False, compare=False)
+    
+    def __len__(self) -> int:
+        """Return number of detections."""
+        return len(self.xyxy)
+    
+    def __getitem__(self, index) -> 'Detections':
+        """Support indexing and slicing."""
+        if isinstance(index, (int, slice, list, np.ndarray)):
+            return Detections(
+                xyxy=self.xyxy[index],
+                confidence=self.confidence[index] if self.confidence is not None else None,
+                class_id=self.class_id[index] if self.class_id is not None else None,
+                tracker_id=self.tracker_id[index] if self.tracker_id is not None else None,
+                mask=self.mask[index] if self.mask is not None else None,
+                data={k: v[index] if isinstance(v, np.ndarray) else v for k, v in self.data.items()},
+            )
+        raise TypeError(f"Invalid index type: {type(index)}")
+    
+    @property
+    def area(self) -> np.ndarray:
+        """Get area of each bounding box (cached)."""
+        if self._area_cache is None:
+            widths = self.xyxy[:, 2] - self.xyxy[:, 0]
+            heights = self.xyxy[:, 3] - self.xyxy[:, 1]
+            self._area_cache = widths * heights
+        return self._area_cache
+    
+    @property
+    def center(self) -> np.ndarray:
+        """Get center point of each bounding box (cached). Shape: (N, 2)"""
+        if self._center_cache is None:
+            self._center_cache = np.column_stack([
+                (self.xyxy[:, 0] + self.xyxy[:, 2]) / 2,
+                (self.xyxy[:, 1] + self.xyxy[:, 3]) / 2
+            ])
+        return self._center_cache
+    
+    @property
+    def width(self) -> np.ndarray:
+        """Get width of each bounding box."""
+        return self.xyxy[:, 2] - self.xyxy[:, 0]
+    
+    @property
+    def height(self) -> np.ndarray:
+        """Get height of each bounding box."""
+        return self.xyxy[:, 3] - self.xyxy[:, 1]
+    
+    @property
+    def aspect_ratio(self) -> np.ndarray:
+        """Get aspect ratio (width/height) of each box."""
+        return self.width / np.maximum(self.height, 1e-6)
+    
+    def filter(self, mask: np.ndarray) -> 'Detections':
+        """Filter detections using boolean mask."""
+        return self[mask]
+    
+    def with_confidence(self, confidence: np.ndarray) -> 'Detections':
+        """Return new Detections with updated confidence."""
+        return Detections(
+            xyxy=self.xyxy,
+            confidence=confidence,
+            class_id=self.class_id,
+            tracker_id=self.tracker_id,
+            mask=self.mask,
+            data=self.data.copy()
+        )
+    
+    def with_class_id(self, class_id: np.ndarray) -> 'Detections':
+        """Return new Detections with updated class IDs."""
+        return Detections(
+            xyxy=self.xyxy,
+            confidence=self.confidence,
+            class_id=class_id,
+            tracker_id=self.tracker_id,
+            mask=self.mask,
+            data=self.data.copy()
+        )
+    
+    def with_tracker_id(self, tracker_id: np.ndarray) -> 'Detections':
+        """Return new Detections with updated tracker IDs."""
+        return Detections(
+            xyxy=self.xyxy,
+            confidence=self.confidence,
+            class_id=self.class_id,
+            tracker_id=tracker_id,
+            mask=self.mask,
+            data=self.data.copy()
+        )
+    
+    @staticmethod
+    def empty() -> 'Detections':
+        """Create empty Detections object."""
+        return Detections(
+            xyxy=np.empty((0, 4), dtype=np.float32),
+            confidence=np.empty(0, dtype=np.float32),
+            class_id=np.empty(0, dtype=int),
+            tracker_id=np.empty(0, dtype=int)
+        )
+    
+    @classmethod
+    def from_yolo(cls, results) -> 'Detections':
+        """Create from YOLO/Ultralytics results."""
+        if hasattr(results, 'boxes'):
+            boxes = results.boxes
+            xyxy = boxes.xyxy.cpu().numpy()
+            confidence = boxes.conf.cpu().numpy()
+            class_id = boxes.cls.cpu().numpy().astype(int)
+            
+            return cls(
+                xyxy=xyxy,
+                confidence=confidence,
+                class_id=class_id
+            )
+        return cls.empty()
+    
+    def merge(self, other: 'Detections') -> 'Detections':
+        """Merge with another Detections object."""
+        return Detections(
+            xyxy=np.vstack([self.xyxy, other.xyxy]),
+            confidence=np.concatenate([self.confidence, other.confidence]) if self.confidence is not None else None,
+            class_id=np.concatenate([self.class_id, other.class_id]) if self.class_id is not None else None,
+            tracker_id=np.concatenate([self.tracker_id, other.tracker_id]) if self.tracker_id is not None else None,
+        )
+
+
+class DetectionManager:
+    """Manages detection history and counting across zones.
+    
+    Innovation: Thread-safe, supports multiple metrics, fast lookups.
+    """
+    
+    def __init__(self, class_names: Optional[Dict[int, str]] = None):
+        self.class_names = class_names or {}
+        self.tracker_to_zone: Dict[int, Tuple[int, int]] = {}  # tracker_id -> (zone_id, class_id)
+        self.tracker_to_metadata: Dict[int, Dict[str, Any]] = {}
+        self.zone_counts: Dict[int, int] = {}
+        
+    def update(
+        self, 
+        detections: Detections, 
+        zone_detections: List[Detections],
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> Detections:
+        """Update tracking information with zone crossings."""
+        for zone_id, zone_det in enumerate(zone_detections):
+            if len(zone_det) == 0:
+                continue
+                
+            for tracker_id, class_id in zip(zone_det.tracker_id, zone_det.class_id):
+                if tracker_id not in self.tracker_to_zone:
+                    self.tracker_to_zone[tracker_id] = (zone_id, class_id)
+                    self.zone_counts[zone_id] = self.zone_counts.get(zone_id, 0) + 1
+                    
+                    if metadata:
+                        self.tracker_to_metadata[tracker_id] = metadata.copy()
+        
+        # Update class IDs based on zone assignment
+        if len(detections) > 0 and detections.tracker_id is not None:
+            new_class_ids = np.array([
+                self.tracker_to_zone.get(tid, (-1, -1))[1] 
+                for tid in detections.tracker_id
+            ])
+            detections = detections.with_class_id(new_class_ids)
+            detections = detections.filter(new_class_ids != -1)
+        
+        return detections
+    
+    def get_zone_count(self, zone_id: int) -> int:
+        """Get count for specific zone."""
+        return self.zone_counts.get(zone_id, 0)
+    
+    def get_tracker_info(self, tracker_id: int) -> Optional[Tuple[int, int, Dict]]:
+        """Get zone, class, and metadata for tracker."""
+        if tracker_id in self.tracker_to_zone:
+            zone_id, class_id = self.tracker_to_zone[tracker_id]
+            metadata = self.tracker_to_metadata.get(tracker_id, {})
+            return zone_id, class_id, metadata
+        return None
+    
+    def export_csv(self, filepath: str):
+        """Export tracking data to CSV."""
+        import csv
+        with open(filepath, 'w', newline='') as f:
+            writer = csv.writer(f)
+            writer.writerow(['tracker_id', 'zone_id', 'class_id', 'class_name', *self.tracker_to_metadata.get(list(self.tracker_to_metadata.keys())[0], {}).keys()])
+            
+            for tracker_id, (zone_id, class_id) in self.tracker_to_zone.items():
+                class_name = self.class_names.get(class_id, str(class_id))
+                metadata = self.tracker_to_metadata.get(tracker_id, {})
+                writer.writerow([tracker_id, zone_id, class_id, class_name, *metadata.values()])
+    
+    def clear(self):
+        """Clear all tracking history."""
+        self.tracker_to_zone.clear()
+        self.tracker_to_metadata.clear()
+        self.zone_counts.clear()
+
+
+# Universal auto-conversion methods
+def from_yolo(results: Any) -> Detections:
+    """Convert YOLO (Ultralytics) results to Detections.
+    
+    Args:
+        results: YOLO results object or list
+    
+    Returns:
+        Detections
+    
+    Example:
+        >>> from ultralytics import YOLO
+        >>> model = YOLO("yolo11n.pt")
+        >>> results = model("image.jpg")
+        >>> detections = eye.from_yolo(results)
+    """
+    # Handle list
+    if isinstance(results, list):
+        if len(results) == 0:
+            return Detections.empty()
+        results = results[0]
+    
+    # Extract boxes
+    boxes = results.boxes.xyxy.cpu().numpy()
+    conf = results.boxes.conf.cpu().numpy() if results.boxes.conf is not None else None
+    cls = results.boxes.cls.cpu().numpy().astype(int) if results.boxes.cls is not None else None
+    
+    # Extract masks if available
+    masks = None
+    if hasattr(results, 'masks') and results.masks is not None:
+        masks = results.masks.data.cpu().numpy()
+    
+    det = Detections(xyxy=boxes, confidence=conf, class_id=cls)
+    if masks is not None:
+        det.data['masks'] = masks
+    
+    return det
+
+
+def from_pytorch(results: Union[Dict, List[Dict]], conf_threshold: float = 0.0) -> Detections:
+    """Convert PyTorch/torchvision results to Detections.
+    
+    Args:
+        results: Dict with 'boxes', 'scores', 'labels' or list of dicts
+        conf_threshold: Confidence threshold
+    
+    Returns:
+        Detections
+    
+    Example:
+        >>> import torchvision
+        >>> model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
+        >>> results = model(image)
+        >>> detections = eye.from_pytorch(results[0])
+    """
+    # Handle list
+    if isinstance(results, list):
+        if len(results) == 0:
+            return Detections.empty()
+        results = results[0]
+    
+    # Extract
+    boxes = results['boxes'].cpu().numpy() if hasattr(results['boxes'], 'cpu') else results['boxes']
+    scores = results['scores'].cpu().numpy() if hasattr(results['scores'], 'cpu') else results['scores']
+    labels = results['labels'].cpu().numpy() if hasattr(results['labels'], 'cpu') else results['labels']
+    
+    # Filter by confidence
+    mask = scores >= conf_threshold
+    
+    return Detections(
+        xyxy=boxes[mask],
+        confidence=scores[mask],
+        class_id=labels[mask].astype(int)
+    )
+
+
+def from_tensorflow(boxes: np.ndarray, scores: np.ndarray, classes: np.ndarray, 
+                   conf_threshold: float = 0.0) -> Detections:
+    """Convert TensorFlow Object Detection API results to Detections.
+    
+    Args:
+        boxes: Detection boxes (normalized [0-1] or absolute)
+        scores: Detection scores
+        classes: Detection classes
+        conf_threshold: Confidence threshold
+    
+    Returns:
+        Detections
+    
+    Example:
+        >>> boxes, scores, classes, num = model(image)
+        >>> detections = eye.from_tensorflow(boxes[0], scores[0], classes[0])
+    """
+    # Filter by confidence
+    mask = scores >= conf_threshold
+    boxes = boxes[mask]
+    scores = scores[mask]
+    classes = classes[mask].astype(int)
+    
+    # Convert normalized to absolute if needed
+    if boxes.max() <= 1.0:
+        print("Warning: Boxes appear normalized. Consider denormalizing.")
+    
+    return Detections(xyxy=boxes, confidence=scores, class_id=classes)
+
+
+def from_opencv(results: np.ndarray, conf_threshold: float = 0.0) -> Detections:
+    """Convert OpenCV DNN results to Detections.
+    
+    Args:
+        results: OpenCV detection output (shape: (1, 1, N, 7))
+        conf_threshold: Confidence threshold
+    
+    Returns:
+        Detections
+    
+    Example:
+        >>> net = cv2.dnn.readNet("model.caffemodel", "deploy.prototxt")
+        >>> blob = cv2.dnn.blobFromImage(image)
+        >>> net.setInput(blob)
+        >>> results = net.forward()
+        >>> detections = eye.from_opencv(results)
+    """
+    # Reshape to (N, 7)
+    if results.ndim == 4:
+        results = results[0, 0]
+    
+    # Filter by confidence
+    mask = results[:, 2] >= conf_threshold
+    filtered = results[mask]
+    
+    # Extract [image_id, label, confidence, x1, y1, x2, y2]
+    boxes = filtered[:, 3:7]
+    scores = filtered[:, 2]
+    classes = filtered[:, 1].astype(int)
+    
+    return Detections(xyxy=boxes, confidence=scores, class_id=classes)
+
+
+def from_numpy(boxes: np.ndarray, scores: Optional[np.ndarray] = None,
+               classes: Optional[np.ndarray] = None) -> Detections:
+    """Convert raw numpy arrays to Detections.
+    
+    Args:
+        boxes: Bounding boxes in xyxy format
+        scores: Confidence scores (optional)
+        classes: Class IDs (optional)
+    
+    Returns:
+        Detections
+    
+    Example:
+        >>> boxes = np.array([[100, 50, 200, 150], [300, 200, 400, 300]])
+        >>> scores = np.array([0.9, 0.85])
+        >>> classes = np.array([0, 1])
+        >>> detections = eye.from_numpy(boxes, scores, classes)
+    """
+    return Detections(xyxy=boxes, confidence=scores, class_id=classes)
+
+
+# Add class methods to Detections
+Detections.from_yolo = staticmethod(from_yolo)
+Detections.from_pytorch = staticmethod(from_pytorch)
+Detections.from_tensorflow = staticmethod(from_tensorflow)
+Detections.from_opencv = staticmethod(from_opencv)
+Detections.from_numpy = staticmethod(from_numpy)