eye-cv 1.0.0__py3-none-any.whl

eye/annotators/trace.py

@@ -0,0 +1,142 @@
+"""Trace annotator for drawing object paths."""
+
+import cv2
+import numpy as np
+from typing import Optional, Dict
+from collections import deque
+from ..core.detections import Detections
+from ..core.colors import ColorPalette
+
+
+class TraceAnnotator:
+    """Draw tracking trails/paths for objects.
+    
+    Innovation: Fading trails, variable thickness, trail smoothing.
+    """
+    
+    def __init__(
+        self,
+        color_palette: Optional[ColorPalette] = None,
+        thickness: int = 2,
+        trace_length: int = 50,
+        fade_trail: bool = True,
+        smooth_trail: bool = False
+    ):
+        """
+        Args:
+            color_palette: Color palette
+            thickness: Line thickness
+            trace_length: Maximum number of points in trail
+            fade_trail: Gradually fade older trail points
+            smooth_trail: Apply smoothing to trail
+        """
+        from ..core.colors import PredefinedPalettes
+        self.color_palette = color_palette or PredefinedPalettes.bright()
+        self.thickness = thickness
+        self.trace_length = trace_length
+        self.fade_trail = fade_trail
+        self.smooth_trail = smooth_trail
+        
+        # Track history for each object
+        self.history: Dict[int, deque] = {}
+        # Position smoothing: base alpha (0-1), lower = more smoothing
+        self.position_smoothing_alpha = 0.6
+        # If movement speed (pixels/frame) exceeds this, use stronger smoothing
+        self.position_speed_threshold = 8.0
+        # Alpha to use when speed > threshold
+        self.fast_position_smoothing_alpha = 0.3
+    
+    def annotate(
+        self,
+        image: np.ndarray,
+        detections: Detections
+    ) -> np.ndarray:
+        """Draw trails on image."""
+        annotated = image.copy()
+        
+        if detections.tracker_id is None:
+            return annotated
+        
+        # Update history
+        active_ids = set()
+        for i in range(len(detections)):
+            tracker_id = detections.tracker_id[i]
+            active_ids.add(tracker_id)
+            center = detections.center[i]
+            
+            if tracker_id not in self.history:
+                self.history[tracker_id] = deque(maxlen=self.trace_length)
+            # Apply optional temporal smoothing when appending to history
+            if len(self.history[tracker_id]) > 0 and self.position_smoothing_alpha is not None:
+                last = np.array(self.history[tracker_id][-1], dtype=float)
+                newc = np.array(center, dtype=float)
+                dist = float(np.linalg.norm(newc - last))
+                if self.position_speed_threshold is not None and dist > float(self.position_speed_threshold):
+                    alpha = float(self.fast_position_smoothing_alpha)
+                else:
+                    alpha = float(self.position_smoothing_alpha)
+                center = alpha * newc + (1.0 - alpha) * last
+
+            self.history[tracker_id].append(center)
+        
+        # Remove old tracks
+        inactive_ids = set(self.history.keys()) - active_ids
+        for tid in inactive_ids:
+            if len(self.history[tid]) > 0:
+                # Keep trail visible for a bit after object leaves
+                pass
+            else:
+                del self.history[tid]
+        
+        # Draw trails
+        for i in range(len(detections)):
+            tracker_id = detections.tracker_id[i]
+            if tracker_id not in self.history or len(self.history[tracker_id]) < 2:
+                continue
+            
+            points = list(self.history[tracker_id])
+            
+            # Smooth if requested
+            if self.smooth_trail and len(points) >= 3:
+                points = self._smooth_points(points)
+            
+            # Get color
+            if detections.class_id is not None:
+                color = self.color_palette.by_id(detections.class_id[i])
+            else:
+                color = self.color_palette.by_id(tracker_id)
+            
+            # Draw trail
+            for j in range(len(points) - 1):
+                if self.fade_trail:
+                    # Calculate alpha based on position in trail
+                    alpha = (j + 1) / len(points)
+                    faded_color = tuple(int(c * alpha) for c in color.as_bgr())
+                else:
+                    faded_color = color.as_bgr()
+                
+                pt1 = tuple(points[j].astype(int))
+                pt2 = tuple(points[j + 1].astype(int))
+                
+                cv2.line(annotated, pt1, pt2, faded_color, self.thickness)
+        
+        return annotated
+    
+    @staticmethod
+    def _smooth_points(points, window=3):
+        """Apply moving average smoothing to points."""
+        points = np.array(points)
+        smoothed = []
+        for i in range(len(points)):
+            start = max(0, i - window // 2)
+            end = min(len(points), i + window // 2 + 1)
+            smoothed.append(np.mean(points[start:end], axis=0))
+        return smoothed
+    
+    def reset(self):
+        """Clear all tracking history."""
+        self.history.clear()
+
+
+# Alias for easier understanding
+PathAnnotator = TraceAnnotator

eye/annotators/utils.py

@@ -0,0 +1,177 @@
+from enum import Enum
+from typing import Optional, Tuple, Union
+
+import numpy as np
+
+from eye.detection.core import Detections
+from eye.draw.color import Color, ColorPalette
+from eye.geometry.core import Position
+
+
+class ColorLookup(Enum):
+    """
+    Enumeration class to define strategies for mapping colors to annotations.
+
+    This enum supports three different lookup strategies:
+        - `INDEX`: Colors are determined by the index of the detection within the scene.
+        - `CLASS`: Colors are determined by the class label of the detected object.
+        - `TRACK`: Colors are determined by the tracking identifier of the object.
+    """
+
+    INDEX = "index"
+    CLASS = "class"
+    TRACK = "track"
+
+    @classmethod
+    def list(cls):
+        return list(map(lambda c: c.value, cls))
+
+
+def resolve_color_idx(
+    detections: Detections,
+    detection_idx: int,
+    color_lookup: Union[ColorLookup, np.ndarray] = ColorLookup.CLASS,
+) -> int:
+    if detection_idx >= len(detections):
+        raise ValueError(
+            f"Detection index {detection_idx} "
+            f"is out of bounds for detections of length {len(detections)}"
+        )
+
+    if isinstance(color_lookup, np.ndarray):
+        if len(color_lookup) != len(detections):
+            raise ValueError(
+                f"Length of color lookup {len(color_lookup)} "
+                f"does not match length of detections {len(detections)}"
+            )
+        return color_lookup[detection_idx]
+    elif color_lookup == ColorLookup.INDEX:
+        return detection_idx
+    elif color_lookup == ColorLookup.CLASS:
+        if detections.class_id is None:
+            raise ValueError(
+                "Could not resolve color by class because "
+                "Detections do not have class_id. If using an annotator, "
+                "try setting color_lookup to sv.ColorLookup.INDEX or "
+                "sv.ColorLookup.TRACK."
+            )
+        return detections.class_id[detection_idx]
+    elif color_lookup == ColorLookup.TRACK:
+        if detections.tracker_id is None:
+            raise ValueError(
+                "Could not resolve color by track because "
+                "Detections do not have tracker_id. Did you call "
+                "tracker.update_with_detections(...) before annotating?"
+            )
+        return detections.tracker_id[detection_idx]
+
+
+def resolve_text_background_xyxy(
+    center_coordinates: Tuple[int, int],
+    text_wh: Tuple[int, int],
+    position: Position,
+) -> Tuple[int, int, int, int]:
+    center_x, center_y = center_coordinates
+    text_w, text_h = text_wh
+
+    if position == Position.TOP_LEFT:
+        return center_x, center_y - text_h, center_x + text_w, center_y
+    elif position == Position.TOP_RIGHT:
+        return center_x - text_w, center_y - text_h, center_x, center_y
+    elif position == Position.TOP_CENTER:
+        return (
+            center_x - text_w // 2,
+            center_y - text_h,
+            center_x + text_w // 2,
+            center_y,
+        )
+    elif position == Position.CENTER or position == Position.CENTER_OF_MASS:
+        return (
+            center_x - text_w // 2,
+            center_y - text_h // 2,
+            center_x + text_w // 2,
+            center_y + text_h // 2,
+        )
+    elif position == Position.BOTTOM_LEFT:
+        return center_x, center_y, center_x + text_w, center_y + text_h
+    elif position == Position.BOTTOM_RIGHT:
+        return center_x - text_w, center_y, center_x, center_y + text_h
+    elif position == Position.BOTTOM_CENTER:
+        return (
+            center_x - text_w // 2,
+            center_y,
+            center_x + text_w // 2,
+            center_y + text_h,
+        )
+    elif position == Position.CENTER_LEFT:
+        return (
+            center_x - text_w,
+            center_y - text_h // 2,
+            center_x,
+            center_y + text_h // 2,
+        )
+    elif position == Position.CENTER_RIGHT:
+        return (
+            center_x,
+            center_y - text_h // 2,
+            center_x + text_w,
+            center_y + text_h // 2,
+        )
+
+
+def get_color_by_index(color: Union[Color, ColorPalette], idx: int) -> Color:
+    if isinstance(color, ColorPalette):
+        return color.by_id(idx)
+    return color
+
+
+def resolve_color(
+    color: Union[Color, ColorPalette],
+    detections: Detections,
+    detection_idx: int,
+    color_lookup: Union[ColorLookup, np.ndarray] = ColorLookup.CLASS,
+) -> Color:
+    idx = resolve_color_idx(
+        detections=detections,
+        detection_idx=detection_idx,
+        color_lookup=color_lookup,
+    )
+    return get_color_by_index(color=color, idx=idx)
+
+
+class Trace:
+    def __init__(
+        self,
+        max_size: Optional[int] = None,
+        start_frame_id: int = 0,
+        anchor: Position = Position.CENTER,
+    ) -> None:
+        self.current_frame_id = start_frame_id
+        self.max_size = max_size
+        self.anchor = anchor
+
+        self.frame_id = np.array([], dtype=int)
+        self.xy = np.empty((0, 2), dtype=np.float32)
+        self.tracker_id = np.array([], dtype=int)
+
+    def put(self, detections: Detections) -> None:
+        frame_id = np.full(len(detections), self.current_frame_id, dtype=int)
+        self.frame_id = np.concatenate([self.frame_id, frame_id])
+        self.xy = np.concatenate(
+            [self.xy, detections.get_anchors_coordinates(self.anchor)]
+        )
+        self.tracker_id = np.concatenate([self.tracker_id, detections.tracker_id])
+
+        unique_frame_id = np.unique(self.frame_id)
+
+        if 0 < self.max_size < len(unique_frame_id):
+            max_allowed_frame_id = self.current_frame_id - self.max_size + 1
+            filtering_mask = self.frame_id >= max_allowed_frame_id
+            self.frame_id = self.frame_id[filtering_mask]
+            self.xy = self.xy[filtering_mask]
+            self.tracker_id = self.tracker_id[filtering_mask]
+
+        self.current_frame_id += 1
+
+    def get(self, tracker_id: int) -> np.ndarray:
+        return self.xy[self.tracker_id == tracker_id]

eye/assets/__init__.py

@@ -0,0 +1,2 @@
+from eye.assets.downloader import download_assets
+from eye.assets.list import VideoAssets

eye/assets/downloader.py

@@ -0,0 +1,95 @@
+import os
+from hashlib import new as hash_new
+from pathlib import Path
+from shutil import copyfileobj
+from typing import Union
+
+from eye.assets.list import VIDEO_ASSETS, VideoAssets
+
+try:
+    from requests import get
+    from tqdm.auto import tqdm
+except ImportError:
+    raise ValueError(
+        "\n"
+        "Please install requests and tqdm to download assets \n"
+        "or install eye with assets \n"
+        "pip install eye[assets] \n"
+        "\n"
+    )
+
+
+def is_md5_hash_matching(filename: str, original_md5_hash: str) -> bool:
+    """
+    Check if the MD5 hash of a file matches the original hash.
+
+    Parameters:
+        filename (str): The path to the file to be checked as a string.
+        original_md5_hash (str): The original MD5 hash to compare against.
+
+    Returns:
+        bool: True if the hashes match, False otherwise.
+    """
+    if not os.path.exists(filename):
+        return False
+
+    with open(filename, "rb") as file:
+        file_contents = file.read()
+        computed_md5_hash = hash_new(name="MD5")
+        computed_md5_hash.update(file_contents)
+
+    return computed_md5_hash.hexdigest() == original_md5_hash
+
+
+def download_assets(asset_name: Union[VideoAssets, str]) -> str:
+    """
+    Download a specified asset if it doesn't already exist or is corrupted.
+
+    Parameters:
+        asset_name (Union[VideoAssets, str]): The name or type of the asset to be
+            downloaded.
+
+    Returns:
+        str: The filename of the downloaded asset.
+
+    Example:
+        ```python
+        from eye.assets import download_assets, VideoAssets
+
+        download_assets(VideoAssets.VEHICLES)
+        "vehicles.mp4"
+        ```
+    """
+
+    filename = asset_name.value if isinstance(asset_name, VideoAssets) else asset_name
+
+    if not Path(filename).exists() and filename in VIDEO_ASSETS:
+        print(f"Downloading {filename} assets \n")
+        response = get(VIDEO_ASSETS[filename][0], stream=True, allow_redirects=True)
+        response.raise_for_status()
+
+        file_size = int(response.headers.get("Content-Length", 0))
+        folder_path = Path(filename).expanduser().resolve()
+        folder_path.parent.mkdir(parents=True, exist_ok=True)
+
+        with tqdm.wrapattr(
+            response.raw, "read", total=file_size, desc="", colour="#a351fb"
+        ) as raw_resp:
+            with folder_path.open("wb") as file:
+                copyfileobj(raw_resp, file)
+
+    elif Path(filename).exists():
+        if not is_md5_hash_matching(filename, VIDEO_ASSETS[filename][1]):
+            print("File corrupted. Re-downloading... \n")
+            os.remove(filename)
+            return download_assets(filename)
+
+        print(f"{filename} asset download complete. \n")
+
+    else:
+        valid_assets = ", ".join(asset.value for asset in VideoAssets)
+        raise ValueError(
+            f"Invalid asset. It should be one of the following: {valid_assets}."
+        )
+
+    return filename

eye/assets/list.py

@@ -0,0 +1,83 @@
+from enum import Enum
+from typing import Dict, Tuple
+
+BASE_VIDEO_URL = "https://media.roboflow.com/eye/video-examples/"
+
+
+class VideoAssets(Enum):
+    """
+    Each member of this enum represents a video asset. The value associated with each
+    member is the filename of the video.
+
+    | Enum Member            | Video Filename             | Video URL                                                                             |
+    |------------------------|----------------------------|---------------------------------------------------------------------------------------|
+    | `VEHICLES`             | `vehicles.mp4`             | [Link](https://media.roboflow.com/eye/video-examples/vehicles.mp4)            |
+    | `MILK_BOTTLING_PLANT`  | `milk-bottling-plant.mp4`  | [Link](https://media.roboflow.com/eye/video-examples/milk-bottling-plant.mp4) |
+    | `VEHICLES_2`           | `vehicles-2.mp4`           | [Link](https://media.roboflow.com/eye/video-examples/vehicles-2.mp4)          |
+    | `GROCERY_STORE`        | `grocery-store.mp4`        | [Link](https://media.roboflow.com/eye/video-examples/grocery-store.mp4)       |
+    | `SUBWAY`               | `subway.mp4`               | [Link](https://media.roboflow.com/eye/video-examples/subway.mp4)              |
+    | `MARKET_SQUARE`        | `market-square.mp4`        | [Link](https://media.roboflow.com/eye/video-examples/market-square.mp4)       |
+    | `PEOPLE_WALKING`       | `people-walking.mp4`       | [Link](https://media.roboflow.com/eye/video-examples/people-walking.mp4)      |
+    | `BEACH`                | `beach-1.mp4`              | [Link](https://media.roboflow.com/eye/video-examples/beach-1.mp4)             |
+    | `BASKETBALL`           | `basketball-1.mp4`         | [Link](https://media.roboflow.com/eye/video-examples/basketball-1.mp4)        |
+    | `SKIING`               | `skiing.mp4`               | [Link](https://media.roboflow.com/eye/video-examples/skiing.mp4)              |
+    """  # noqa: E501 // docs
+
+    VEHICLES = "vehicles.mp4"
+    MILK_BOTTLING_PLANT = "milk-bottling-plant.mp4"
+    VEHICLES_2 = "vehicles-2.mp4"
+    GROCERY_STORE = "grocery-store.mp4"
+    SUBWAY = "subway.mp4"
+    MARKET_SQUARE = "market-square.mp4"
+    PEOPLE_WALKING = "people-walking.mp4"
+    BEACH = "beach-1.mp4"
+    BASKETBALL = "basketball-1.mp4"
+    SKIING = "skiing.mp4"
+
+    @classmethod
+    def list(cls):
+        return list(map(lambda c: c.value, cls))
+
+
+VIDEO_ASSETS: Dict[str, Tuple[str, str]] = {
+    VideoAssets.VEHICLES.value: (
+        f"{BASE_VIDEO_URL}{VideoAssets.VEHICLES.value}",
+        "8155ff4e4de08cfa25f39de96483f918",
+    ),
+    VideoAssets.VEHICLES_2.value: (
+        f"{BASE_VIDEO_URL}{VideoAssets.VEHICLES_2.value}",
+        "830af6fba21ffbf14867a7fea595937b",
+    ),
+    VideoAssets.MILK_BOTTLING_PLANT.value: (
+        f"{BASE_VIDEO_URL}{VideoAssets.MILK_BOTTLING_PLANT.value}",
+        "9e8fb6e883f842a38b3d34267290bdc7",
+    ),
+    VideoAssets.GROCERY_STORE.value: (
+        f"{BASE_VIDEO_URL}{VideoAssets.GROCERY_STORE.value}",
+        "11402e7b861c1980527d3d74cbe3b366",
+    ),
+    VideoAssets.SUBWAY.value: (
+        f"{BASE_VIDEO_URL}{VideoAssets.SUBWAY.value}",
+        "453475750691fb23c56a0cffef089194",
+    ),
+    VideoAssets.MARKET_SQUARE.value: (
+        f"{BASE_VIDEO_URL}{VideoAssets.MARKET_SQUARE.value}",
+        "859179bf4a21f80a8baabfdb2ed716dc",
+    ),
+    VideoAssets.PEOPLE_WALKING.value: (
+        f"{BASE_VIDEO_URL}{VideoAssets.PEOPLE_WALKING.value}",
+        "0574c053c8686c3f1dc0aa3743e45cb9",
+    ),
+    VideoAssets.BEACH.value: (
+        f"{BASE_VIDEO_URL}{VideoAssets.BEACH.value}",
+        "4175d42fec4d450ed081523fd39e0cf8",
+    ),
+    VideoAssets.BASKETBALL.value: (
+        f"{BASE_VIDEO_URL}{VideoAssets.BASKETBALL.value}",
+        "60d94a3c7c47d16f09d342b088012ecc",
+    ),
+    VideoAssets.SKIING.value: (
+        f"{BASE_VIDEO_URL}{VideoAssets.SKIING.value}",
+        "d30987cbab1bbc5934199cdd1b293119",
+    ),
+}

eye/classification/core.py

@@ -0,0 +1,188 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Optional, Tuple
+
+import numpy as np
+
+
+def _validate_class_ids(class_id: Any, n: int) -> None:
+    """
+    Ensure that class_id is a 1d np.ndarray with (n, ) shape.
+    """
+    is_valid = isinstance(class_id, np.ndarray) and class_id.shape == (n,)
+    if not is_valid:
+        raise ValueError("class_id must be 1d np.ndarray with (n, ) shape")
+
+
+def _validate_confidence(confidence: Any, n: int) -> None:
+    """
+    Ensure that confidence is a 1d np.ndarray with (n, ) shape.
+    """
+    if confidence is not None:
+        is_valid = isinstance(confidence, np.ndarray) and confidence.shape == (n,)
+        if not is_valid:
+            raise ValueError("confidence must be 1d np.ndarray with (n, ) shape")
+
+
+@dataclass
+class Classifications:
+    class_id: np.ndarray
+    confidence: Optional[np.ndarray] = None
+
+    def __post_init__(self) -> None:
+        """
+        Validate the classification inputs.
+        """
+        n = len(self.class_id)
+
+        _validate_class_ids(self.class_id, n)
+        _validate_confidence(self.confidence, n)
+
+    def __len__(self) -> int:
+        """
+        Returns the number of classifications.
+        """
+        return len(self.class_id)
+
+    @classmethod
+    def from_clip(cls, clip_results) -> Classifications:
+        """
+        Creates a Classifications instance from a
+        [clip](https://github.com/openai/clip) inference result.
+
+        Args:
+            clip_results (np.ndarray): The inference result from clip model.
+
+        Returns:
+            Classifications: A new Classifications object.
+
+        Example:
+            ```python
+            from PIL import Image
+            import clip
+            import eye as sv
+
+            model, preprocess = clip.load('ViT-B/32')
+
+            image = cv2.imread(SOURCE_IMAGE_PATH)
+            image = preprocess(image).unsqueeze(0)
+
+            text = clip.tokenize(["a diagram", "a dog", "a cat"])
+            output, _ = model(image, text)
+            classifications = sv.Classifications.from_clip(output)
+            ```
+        """
+
+        confidence = clip_results.softmax(dim=-1).cpu().detach().numpy()[0]
+
+        if len(confidence) == 0:
+            return cls(class_id=np.array([]), confidence=np.array([]))
+
+        class_ids = np.arange(len(confidence))
+        return cls(class_id=class_ids, confidence=confidence)
+
+    @classmethod
+    def from_ultralytics(cls, ultralytics_results) -> Classifications:
+        """
+        Creates a Classifications instance from a
+        [ultralytics](https://github.com/ultralytics/ultralytics) inference result.
+
+        Args:
+            ultralytics_results (ultralytics.engine.results.Results):
+                The inference result from ultralytics model.
+
+        Returns:
+            Classifications: A new Classifications object.
+
+        Example:
+            ```python
+            import cv2
+            from ultralytics import YOLO
+            import eye as sv
+
+            image = cv2.imread(SOURCE_IMAGE_PATH)
+            model = YOLO('yolov8n-cls.pt')
+
+            output = model(image)[0]
+            classifications = sv.Classifications.from_ultralytics(output)
+            ```
+        """
+        confidence = ultralytics_results.probs.data.cpu().numpy()
+        return cls(class_id=np.arange(confidence.shape[0]), confidence=confidence)
+
+    @classmethod
+    def from_timm(cls, timm_results) -> Classifications:
+        """
+        Creates a Classifications instance from a
+        [timm](https://huggingface.co/docs/hub/timm) inference result.
+
+        Args:
+            timm_results (torch.Tensor): The inference result from timm model.
+
+        Returns:
+            Classifications: A new Classifications object.
+
+        Example:
+            ```python
+            from PIL import Image
+            import timm
+            from timm.data import resolve_data_config, create_transform
+            import eye as sv
+
+            model = timm.create_model(
+                model_name='hf-hub:nateraw/resnet50-oxford-iiit-pet',
+                pretrained=True
+            ).eval()
+
+            config = resolve_data_config({}, model=model)
+            transform = create_transform(**config)
+
+            image = Image.open(SOURCE_IMAGE_PATH).convert('RGB')
+            x = transform(image).unsqueeze(0)
+
+            output = model(x)
+
+            classifications = sv.Classifications.from_timm(output)
+            ```
+        """
+        confidence = timm_results.cpu().detach().numpy()[0]
+
+        if len(confidence) == 0:
+            return cls(class_id=np.array([]), confidence=np.array([]))
+
+        class_id = np.arange(len(confidence))
+        return cls(class_id=class_id, confidence=confidence)
+
+    def get_top_k(self, k: int) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Retrieve the top k class IDs and confidences,
+            ordered in descending order by confidence.
+
+        Args:
+            k (int): The number of top class IDs and confidences to retrieve.
+
+        Returns:
+            Tuple[np.ndarray, np.ndarray]: A tuple containing
+                the top k class IDs and confidences.
+
+        Example:
+            ```python
+            import eye as sv
+
+            classifications = sv.Classifications(...)
+
+            classifications.get_top_k(1)
+
+            (array([1]), array([0.9]))
+            ```
+        """
+        if self.confidence is None:
+            raise ValueError("top_k could not be calculated, confidence is None")
+
+        order = np.argsort(self.confidence)[::-1]
+        top_k_order = order[:k]
+        top_k_class_id = self.class_id[top_k_order]
+        top_k_confidence = self.confidence[top_k_order]
+
+        return top_k_class_id, top_k_confidence

eye/config.py

@@ -0,0 +1,2 @@
+CLASS_NAME_DATA_FIELD = "class_name"
+ORIENTED_BOX_COORDINATES = "xyxyxyxy"