superiorvision 0.30.0.dev0__py3-none-any.whl

supervision/assets/list.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+from enum import Enum
+
+BASE_VIDEO_URL = "https://media.roboflow.com/supervision/video-examples/"
+BASE_IMAGE_URL = "https://media.roboflow.com/supervision/image-examples/"
+
+
+class Assets(Enum):
+    filename: str
+    md5_hash: str
+
+    def __new__(cls, filename: str, md5_hash: str) -> Assets:
+        obj = object.__new__(cls)
+        obj._value_ = filename
+        obj.filename = filename
+        obj.md5_hash = md5_hash
+        return obj
+
+    @classmethod
+    def list(cls) -> list[str]:
+        return [asset.filename for asset in cls]
+
+
+class VideoAssets(Assets):
+    """
+    Each member of this class represents a video asset. The value associated with each
+    member has a filename and hash of the video. File names and links can be seen below.
+
+    | Asset                  | Video Filename             | Video URL                                                                             |
+    |------------------------|----------------------------|---------------------------------------------------------------------------------------|
+    | `VEHICLES`             | `vehicles.mp4`             | [Link](https://media.roboflow.com/supervision/video-examples/vehicles.mp4)            |
+    | `MILK_BOTTLING_PLANT`  | `milk-bottling-plant.mp4`  | [Link](https://media.roboflow.com/supervision/video-examples/milk-bottling-plant.mp4) |
+    | `VEHICLES_2`           | `vehicles-2.mp4`           | [Link](https://media.roboflow.com/supervision/video-examples/vehicles-2.mp4)          |
+    | `GROCERY_STORE`        | `grocery-store.mp4`        | [Link](https://media.roboflow.com/supervision/video-examples/grocery-store.mp4)       |
+    | `SUBWAY`               | `subway.mp4`               | [Link](https://media.roboflow.com/supervision/video-examples/subway.mp4)              |
+    | `MARKET_SQUARE`        | `market-square.mp4`        | [Link](https://media.roboflow.com/supervision/video-examples/market-square.mp4)       |
+    | `PEOPLE_WALKING`       | `people-walking.mp4`       | [Link](https://media.roboflow.com/supervision/video-examples/people-walking.mp4)      |
+    | `BEACH`                | `beach-1.mp4`              | [Link](https://media.roboflow.com/supervision/video-examples/beach-1.mp4)             |
+    | `BASKETBALL`           | `basketball-1.mp4`         | [Link](https://media.roboflow.com/supervision/video-examples/basketball-1.mp4)        |
+    | `SKIING`               | `skiing.mp4`               | [Link](https://media.roboflow.com/supervision/video-examples/skiing.mp4)              |
+    """  # noqa: E501 // docs
+
+    VEHICLES = ("vehicles.mp4", "8155ff4e4de08cfa25f39de96483f918")
+    MILK_BOTTLING_PLANT = (
+        "milk-bottling-plant.mp4",
+        "9e8fb6e883f842a38b3d34267290bdc7",
+    )
+    VEHICLES_2 = ("vehicles-2.mp4", "830af6fba21ffbf14867a7fea595937b")
+    GROCERY_STORE = ("grocery-store.mp4", "48608fb4a8981f1c2469fa492adeec9c")
+    SUBWAY = ("subway.mp4", "453475750691fb23c56a0cffef089194")
+    MARKET_SQUARE = ("market-square.mp4", "859179bf4a21f80a8baabfdb2ed716dc")
+    PEOPLE_WALKING = ("people-walking.mp4", "0574c053c8686c3f1dc0aa3743e45cb9")
+    BEACH = ("beach-1.mp4", "4175d42fec4d450ed081523fd39e0cf8")
+    BASKETBALL = ("basketball-1.mp4", "60d94a3c7c47d16f09d342b088012ecc")
+    SKIING = ("skiing.mp4", "d30987cbab1bbc5934199cdd1b293119")
+
+
+class ImageAssets(Assets):
+    """
+    Each member of this enum represents a image asset. The value associated with each
+    member is the filename of the image.
+
+    | Asset              | Image Filename         | Video URL                                                                             |
+    |--------------------|------------------------|---------------------------------------------------------------------------------------|
+    | `PEOPLE_WALKING`   | `people-walking.jpg`   | [Link](https://media.roboflow.com/supervision/image-examples/people-walking.jpg)      |
+    | `SOCCER`           | `soccer.jpg`           | [Link](https://media.roboflow.com/supervision/image-examples/soccer.jpg)              |
+
+    """  # noqa: E501 // docs
+
+    PEOPLE_WALKING = ("people-walking.jpg", "e6bda00b47f2908eeae7df86ef995dcd")
+    SOCCER = ("soccer.jpg", "0f5a4b98abf3e3973faf9e9260a7d876")
+
+
+MEDIA_ASSETS: dict[str, tuple[str, str]] = {
+    **{
+        asset.filename: (f"{BASE_VIDEO_URL}{asset.filename}", asset.md5_hash)
+        for asset in VideoAssets
+    },
+    **{
+        asset.filename: (f"{BASE_IMAGE_URL}{asset.filename}", asset.md5_hash)
+        for asset in ImageAssets
+    },
+}

supervision/classification/core.py

@@ -0,0 +1,216 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import numpy.typing as npt
+
+if TYPE_CHECKING:
+    import torch  # type: ignore[import-not-found, unused-ignore]
+
+
+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: npt.NDArray[np.int_]
+    confidence: npt.NDArray[np.floating] | None = 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 __eq__(self, other: object) -> bool:
+        """
+        Compare classifications by value across numpy-backed fields.
+        """
+        if not isinstance(other, Classifications):
+            return NotImplemented
+        if not np.array_equal(self.class_id, other.class_id):
+            return False
+        if self.confidence is None or other.confidence is None:
+            return self.confidence is other.confidence
+        return bool(np.array_equal(self.confidence, other.confidence))
+
+    def __len__(self) -> int:
+        """
+        Returns the number of classifications.
+        """
+        return len(self.class_id)
+
+    @classmethod
+    def from_clip(cls, clip_results: torch.Tensor) -> Classifications:
+        """
+        Creates a Classifications instance from a
+        [clip](https://github.com/openai/clip) inference result.
+
+        Args:
+            clip_results: The inference result from clip model.
+
+        Returns:
+            A new Classifications object.
+
+        Example:
+            ```python
+            from PIL import Image
+            import clip
+            import supervision 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([], dtype=np.int_),
+                confidence=np.array([], dtype=np.float32),
+            )
+
+        class_ids = np.arange(len(confidence))
+        return cls(class_id=class_ids, confidence=confidence)
+
+    @classmethod
+    def from_ultralytics(cls, ultralytics_results: Any) -> Classifications:
+        """
+        Creates a Classifications instance from a
+        [ultralytics](https://github.com/ultralytics/ultralytics) inference result.
+
+        Args:
+            ultralytics_results: The inference result from ultralytics model.
+
+        Returns:
+            A new Classifications object.
+
+        Example:
+            ```python
+            import cv2
+            from ultralytics import YOLO
+            import supervision 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: Any) -> Classifications:
+        """
+        Creates a Classifications instance from a
+        [timm](https://huggingface.co/docs/hub/timm) inference result.
+
+        Note:
+            Returned confidences are softmax-normalized probabilities, so
+            thresholds calibrated against raw logits may need recalibration.
+
+        Args:
+            timm_results: The inference result from timm model.
+
+        Returns:
+            A new Classifications object.
+
+        Example:
+            ```python
+            from PIL import Image
+            import timm
+            from timm.data import resolve_data_config, create_transform
+            import supervision 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.softmax(dim=-1).cpu().detach().numpy()[0]
+
+        if len(confidence) == 0:
+            return cls(
+                class_id=np.array([], dtype=np.int_),
+                confidence=np.array([], dtype=np.float32),
+            )
+
+        class_id = np.arange(len(confidence))
+        return cls(class_id=class_id, confidence=confidence)
+
+    def get_top_k(
+        self, k: int
+    ) -> tuple[npt.NDArray[np.int_], npt.NDArray[np.floating]]:
+        """
+        Retrieve the top k class IDs and confidences,
+            ordered in descending order by confidence.
+
+        Args:
+            k: The number of top class IDs and confidences to retrieve.
+
+        Returns:
+            A tuple containing the top k class IDs and confidences.
+
+        Example:
+            ```pycon
+            >>> import numpy as np
+            >>> import supervision as sv
+            >>> classifications = sv.Classifications(
+            ...     class_id=np.array([0, 1, 2]),
+            ...     confidence=np.array([0.3, 0.9, 0.5])
+            ... )
+            >>> 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

supervision/config.py

@@ -0,0 +1,13 @@
+CLASS_NAME_DATA_FIELD: str = "class_name"
+COCO_RAW_SEGMENTATION: str = "coco_raw_segmentation"
+#: Key for oriented bounding-box corner coordinates in ``Detections.data``.
+#:
+#: Value layout: ``np.ndarray`` of shape ``(N, 4, 2)``, dtype ``float32``, pixel
+#: coordinates ordered as ``[[x1, y1], [x2, y2], [x3, y3], [x4, y4]]`` per
+#: detection where the four points are the corners of the oriented box.
+#: Used by :func:`~supervision.dataset.formats.yolo.detections_to_yolo_annotations`
+#: (``is_obb=True``) and
+#: :func:`~supervision.dataset.formats.yolo.yolo_annotations_to_detections`
+#: (``is_obb=True``).
+#: Also triggers sequential mode in ``InferenceSlicer`` when present.
+ORIENTED_BOX_COORDINATES: str = "xyxyxyxy"