superiorvision 0.30.0.dev0__py3-none-any.whl

supervision/annotators/utils.py

@@ -0,0 +1,541 @@
+import re
+import textwrap
+from enum import Enum
+from typing import Any, cast
+
+import numpy as np
+import numpy.typing as npt
+
+from supervision.config import CLASS_NAME_DATA_FIELD
+from supervision.detection.core import Detections
+from supervision.draw.color import Color, ColorPalette
+from supervision.geometry.core import Position
+from supervision.utils.deprecate import deprecated, void
+from supervision.utils.tensor import to_numpy
+
+PENDING_TRACK_COLOR = Color.GREY
+PENDING_TRACK_ID = -1
+
+
+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) -> list[str]:
+        return list(map(lambda c: c.value, cls))
+
+
+def resolve_color_idx(
+    detections: Detections,
+    detection_idx: int,
+    color_lookup: ColorLookup | npt.NDArray[np.int_] = 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 int(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 int(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 int(detections.tracker_id[detection_idx])
+    raise ValueError(f"Unsupported color lookup strategy: {color_lookup}")
+
+
+def resolve_text_background_xyxy(
+    center_coordinates: tuple[int, int],
+    text_wh: tuple[int, int],
+    position: Position,
+) -> tuple[int, int, int, int]:
+    """Compute the background box for text anchored at `position`."""
+    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,
+        )
+    raise ValueError(f"Unsupported position: {position}")
+
+
+def get_color_by_index(color: Color | ColorPalette, idx: int) -> Color:
+    """Resolve a color-like object to a concrete `Color` for an index."""
+    color_like = cast(Any, color)
+    # Accept ColorPalette-like objects without depending on their exact concrete class.
+    if callable(getattr(color_like, "by_idx", None)):
+        color_like = color_like.by_idx(idx)
+    if isinstance(color_like, Color):
+        return color_like
+    return Color(
+        r=int(color_like.r),
+        g=int(color_like.g),
+        b=int(color_like.b),
+        a=int(getattr(color_like, "a", 255)),
+    )
+
+
+def resolve_color(
+    color: Color | ColorPalette,
+    detections: Detections,
+    detection_idx: int,
+    color_lookup: ColorLookup | npt.NDArray[np.int_] = ColorLookup.CLASS,
+) -> Color:
+    idx = resolve_color_idx(
+        detections=detections,
+        detection_idx=detection_idx,
+        color_lookup=color_lookup,
+    )
+    if (
+        isinstance(color_lookup, ColorLookup)
+        and color_lookup == ColorLookup.TRACK
+        and idx == PENDING_TRACK_ID
+    ):
+        return PENDING_TRACK_COLOR
+    return get_color_by_index(color=color, idx=idx)
+
+
+def wrap_text(text: object, max_line_length: int | None = None) -> list[str]:
+    """
+    Wrap `text` to the specified maximum line length, respecting existing
+    newlines. Falls back to str() if `text` is not already a string.
+
+    Args:
+        text: The text (or object) to wrap.
+        max_line_length: Maximum width for each wrapped line.
+
+    Returns:
+        Wrapped lines.
+    """
+
+    if not text:
+        return [""]
+
+    if not isinstance(text, str):
+        text = str(text)
+
+    if max_line_length is None:
+        return text.splitlines() or [""]
+
+    if max_line_length <= 0:
+        raise ValueError("max_line_length must be a positive integer")
+
+    paragraphs = text.split("\n")
+    all_lines: list[str] = []
+
+    for paragraph in paragraphs:
+        if paragraph == "":
+            all_lines.append("")
+            continue
+
+        wrapped = textwrap.wrap(
+            paragraph,
+            width=max_line_length,
+            break_long_words=True,
+            replace_whitespace=False,
+            drop_whitespace=True,
+        )
+
+        all_lines.extend(wrapped or [""])
+
+    return all_lines or [""]
+
+
+def _validate_labels(labels: list[str] | None, detections: Detections) -> None:
+    """
+    Validates that the number of provided labels matches the number of detections.
+
+    Args:
+        labels: A list of labels, one for each detection. Can
+            be None.
+        detections: The detections to be labeled.
+
+    Raises:
+        ValueError: If `labels` is not None and its length does not match the number
+            of detections.
+    """
+    if labels is not None and len(labels) != len(detections):
+        raise ValueError(
+            f"The number of labels ({len(labels)}) does not match the "
+            f"number of detections ({len(detections)}). Each detection "
+            f"should have exactly 1 label."
+        )
+
+
+@deprecated(  # type: ignore[untyped-decorator]
+    target=_validate_labels,
+    deprecated_in="0.29.0",
+    remove_in="0.32.0",
+)
+def validate_labels(labels: list[str] | None, detections: Detections) -> None:
+    void(labels, detections)
+
+
+def get_labels_text(
+    detections: Detections, custom_labels: list[str] | None
+) -> list[str]:
+    """
+    Retrieves the text labels for the detections.
+
+    If `custom_labels` are provided, they are used. Otherwise, the labels are
+    extracted from the `detections` object, prioritizing the 'class_name' field,
+    then the `class_id`, and finally using the detection index as a string.
+
+    Args:
+        detections: The detections to get labels for.
+        custom_labels: An optional list of custom labels.
+
+    Returns:
+        A list of text labels for each detection.
+    """
+    if custom_labels is not None:
+        return custom_labels
+
+    labels = []
+    for idx in range(len(detections)):
+        if CLASS_NAME_DATA_FIELD in detections.data:
+            labels.append(str(detections.data[CLASS_NAME_DATA_FIELD][idx]))
+        elif detections.class_id is not None:
+            labels.append(str(int(to_numpy(detections.class_id[idx]))))
+        else:
+            labels.append(str(idx))
+    return labels
+
+
+def snap_boxes(
+    xyxy: npt.NDArray[np.float32],
+    resolution_wh: tuple[int, int],
+) -> npt.NDArray[np.float32]:
+    """
+    Shifts `label` bounding boxes into the frame so that they are fully contained
+    within the given resolution, prioritizing the top/left edge.
+    Unlike `clip_boxes`, this function does not crop boxes.
+    It moves them entirely if they exceed the frame boundaries.
+
+    Args:
+        xyxy: A numpy array of shape `(N, 4)` where each
+            row corresponds to a bounding box in the format
+            `(x_min, y_min, x_max, y_max)`.
+        resolution_wh: A tuple `(width, height)`
+            representing the resolution of the frame.
+
+    Returns:
+        A numpy array of shape `(N, 4)` with boxes shifted into frame.
+
+    Examples:
+        ```pycon
+        >>> import numpy as np
+        >>> from supervision.annotators.utils import snap_boxes
+        >>> xyxy = np.array([
+        ...     [-10, 10, 30, 50],     # Off left edge
+        ...     [310, 200, 350, 250],  # Off right edge
+        ...     [100, -20, 150, 30],   # Off top edge
+        ...     [200, 220, 250, 270],  # Off bottom edge
+        ...     [-20, 10, 350, 50],    # Wider than frame (370 vs 320)
+        ...     [10, -20, 30, 260]     # Taller than frame (280 vs 240)
+        ... ])
+        >>> resolution_wh = (320, 240)
+        >>> snapped_boxes = snap_boxes(xyxy=xyxy, resolution_wh=resolution_wh)
+        >>> snapped_boxes
+        array([[  0.,  10.,  40.,  50.],
+               [280., 190., 320., 240.],
+               [100.,   0., 150.,  50.],
+               [200., 190., 250., 240.],
+               [  0.,  10., 370.,  50.],
+               [ 10.,   0.,  30., 280.]], dtype=float32)
+
+        ```
+    """
+    result: npt.NDArray[np.float32] = np.array(xyxy, dtype=np.float32, copy=True)
+    width, height = resolution_wh
+
+    # X-axis (prioritize left edge)
+    left_overflow = result[:, 0] < 0
+    result[left_overflow, 0:3:2] -= result[left_overflow, 0:1]
+
+    right_overflow = (~left_overflow) & (result[:, 2] > width)
+    right_shift = width - result[right_overflow, 2]
+    result[right_overflow, 0:3:2] += right_shift[:, np.newaxis]
+
+    # Y-axis (prioritize top edge)
+    top_overflow = result[:, 1] < 0
+    result[top_overflow, 1:4:2] -= result[top_overflow, 1:2]
+
+    bottom_overflow = (~top_overflow) & (result[:, 3] > height)
+    bottom_shift = height - result[bottom_overflow, 3]
+    result[bottom_overflow, 1:4:2] += bottom_shift[:, np.newaxis]
+
+    return cast(npt.NDArray[np.float32], result.astype(np.float32, copy=False))
+
+
+class Trace:
+    def __init__(
+        self,
+        max_size: int | None = 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: npt.NDArray[np.int_] = np.array([], dtype=int)
+        self.xy: npt.NDArray[np.float32] = np.empty((0, 2), dtype=np.float32)
+        self.tracker_id: npt.NDArray[np.int_] = np.array([], dtype=int)
+
+    def put(self, detections: Detections) -> None:
+        """Append a frame of detections to the trace history."""
+        if detections.tracker_id is None:
+            raise ValueError(
+                "Could not put detections into Trace because "
+                "Detections do not have tracker_id."
+            )
+
+        frame_id: npt.NDArray[np.int_] = 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,
+                to_numpy(detections.get_anchors_coordinates(self.anchor)),
+            ]
+        )
+        self.tracker_id = np.concatenate(
+            [self.tracker_id, to_numpy(detections.tracker_id)]
+        )
+
+        unique_frame_id = np.unique(self.frame_id)
+
+        if self.max_size is not None and 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) -> npt.NDArray[np.float32]:
+        xy: npt.NDArray[np.float32] = np.asarray(
+            self.xy[self.tracker_id == tracker_id], dtype=np.float32
+        )
+        return xy
+
+    def reset(self) -> None:
+        """Restore the trace buffers to their initial empty state.
+
+        Clears the accumulated `frame_id`, `xy`, and `tracker_id` history and
+        rewinds `current_frame_id` to `0`, so the trace can be reused across
+        independent streams without carrying over points from a previous run.
+        """
+        self.current_frame_id = 0
+        self.frame_id = np.array([], dtype=int)
+        self.xy = np.empty((0, 2), dtype=np.float32)
+        self.tracker_id = np.array([], dtype=int)
+
+
+def hex_to_rgba(hex_color: str) -> tuple[int, int, int, int]:
+    """
+    Converts a hex color string (e.g. "#FF00FF" or "#FF00FF80") to an RGBA tuple.
+
+    Args:
+        hex_color: A hex color string.
+
+    Returns:
+        RGBA values in range 0-255.
+
+    Raises:
+        ValueError: If the format is invalid.
+
+    Examples:
+        ```pycon
+        >>> from supervision.annotators.utils import hex_to_rgba
+        >>> hex_to_rgba("#FF00FF")
+        (255, 0, 255, 255)
+        >>> hex_to_rgba("#FF00FF80")
+        (255, 0, 255, 128)
+
+        ```
+    """
+    hex_color = hex_color.strip().removeprefix("#")
+    if len(hex_color) == 6:
+        hex_color += "FF"  # default full opacity
+    if len(hex_color) != 8:
+        raise ValueError(f"Invalid hex color format: {hex_color}")
+    try:
+        r = int(hex_color[0:2], 16)
+        g = int(hex_color[2:4], 16)
+        b = int(hex_color[4:6], 16)
+        a = int(hex_color[6:8], 16)
+    except ValueError as exc:
+        raise ValueError(f"Invalid hex digits in {hex_color}") from exc
+    return (r, g, b, a)
+
+
+def rgba_to_hex(rgba: tuple[int, int, int, int]) -> str:
+    """
+    Converts an RGBA tuple (0-255 each) to a hex color string.
+
+    Args:
+        rgba: RGBA values in range 0-255.
+
+    Returns:
+        Hex color string in the format "#RRGGBBAA".
+
+    Raises:
+        ValueError: If `rgba` is not a 4-tuple or contains values outside 0-255.
+
+    Examples:
+        ```pycon
+        >>> from supervision.annotators.utils import rgba_to_hex
+        >>> rgba_to_hex((255, 0, 255, 128))
+        '#FF00FF80'
+
+        ```
+    """
+    if len(rgba) != 4 or not all(0 <= c <= 255 for c in rgba):
+        raise ValueError("RGBA must be a 4-tuple with values between 0-255.")
+    return "#{:02X}{:02X}{:02X}{:02X}".format(*rgba)
+
+
+def is_valid_hex(hex_color: str) -> bool:
+    """
+    Checks if a given string is a valid hex color.
+
+    Args:
+        hex_color: A hex color string with an optional leading "#". Supports
+            6-digit (RGB) or 8-digit (RGBA) formats.
+
+    Returns:
+        True if the string is a valid 6- or 8-digit hex color, otherwise False.
+
+    Examples:
+        ```pycon
+        >>> from supervision.annotators.utils import is_valid_hex
+        >>> is_valid_hex("#FF00FF")
+        True
+        >>> is_valid_hex("not-a-color")
+        False
+
+        ```
+    """
+    return bool(re.fullmatch(r"#?[0-9A-Fa-f]{6}([0-9A-Fa-f]{2})?", hex_color.strip()))
+
+
+def calculate_dynamic_kernel_size(x1: int, y1: int, x2: int, y2: int) -> int:
+    """
+    Computes a blur kernel size proportional to the shorter side of a bounding box.
+
+    Args:
+        x1: Left edge of the bounding box.
+        y1: Top edge of the bounding box.
+        x2: Right edge of the bounding box.
+        y2: Bottom edge of the bounding box.
+
+    Returns:
+        Kernel size as one-third of the shorter dimension, minimum 1.
+
+    Examples:
+        ```pycon
+        >>> calculate_dynamic_kernel_size(0, 0, 90, 60)
+        20
+
+        ```
+    """
+    return max(1, min(y2 - y1, x2 - x1) // 3)
+
+
+def calculate_dynamic_pixel_size(x1: int, y1: int, x2: int, y2: int) -> int:
+    """
+    Computes a pixelation size proportional to the shorter side of a bounding box.
+
+    Args:
+        x1: Left edge of the bounding box.
+        y1: Top edge of the bounding box.
+        x2: Right edge of the bounding box.
+        y2: Bottom edge of the bounding box.
+
+    Returns:
+        Pixel size as one-half of the shorter dimension, minimum 1.
+
+    Examples:
+        ```pycon
+        >>> calculate_dynamic_pixel_size(0, 0, 90, 60)
+        30
+
+        ```
+    """
+    return max(1, min(y2 - y1, x2 - x1) // 2)

supervision/assets/__init__.py

@@ -0,0 +1,4 @@
+from supervision.assets.downloader import download_assets
+from supervision.assets.list import ImageAssets, VideoAssets
+
+__all__ = ["ImageAssets", "VideoAssets", "download_assets"]

supervision/assets/downloader.py

@@ -0,0 +1,166 @@
+import os
+from hashlib import md5
+from pathlib import Path
+from shutil import copyfileobj
+
+from requests import get
+from tqdm.auto import tqdm
+
+from supervision.assets.list import MEDIA_ASSETS, Assets
+from supervision.utils.logger import _get_logger
+
+logger = _get_logger(__name__)
+
+
+def is_md5_hash_matching(filename: str | Path, original_md5_hash: str) -> bool:
+    """
+    Check if the MD5 hash of a file matches the original hash.
+
+    Note: MD5 is used here for file integrity checking (detecting corruption),
+    not for cryptographic security purposes.
+
+    Args:
+        filename: The path to the file to be checked.
+        original_md5_hash: The original MD5 hash to compare against.
+
+    Returns:
+        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 = md5(file_contents, usedforsecurity=False)
+
+    return computed_md5_hash.hexdigest() == original_md5_hash
+
+
+def _download_asset(filename: str, destination: Path) -> None:
+    """
+    Download asset bytes to the target destination via a temporary file.
+    """
+    response = get(
+        MEDIA_ASSETS[filename][0], stream=True, allow_redirects=True, timeout=30
+    )
+    response.raise_for_status()
+
+    file_size = int(response.headers.get("Content-Length", 0))
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    temp_path = destination.with_name(f"{destination.name}.part")
+
+    try:
+        with tqdm.wrapattr(
+            response.raw, "read", total=file_size, desc="", colour="#a351fb"
+        ) as raw_resp:
+            with temp_path.open("wb") as file:
+                copyfileobj(raw_resp, file)
+    except Exception:
+        temp_path.unlink(missing_ok=True)
+        raise
+
+    try:
+        os.replace(temp_path, destination)
+    finally:
+        temp_path.unlink(missing_ok=True)
+
+
+def _download_verified_asset(
+    filename: str,
+    original_md5_hash: str,
+    destination: Path,
+    check_target: str | Path,
+    retry_on_mismatch: bool = True,
+) -> None:
+    """
+    Download an asset and reject payloads whose MD5 does not match the catalog.
+    """
+    _download_asset(filename, destination)
+
+    if is_md5_hash_matching(check_target, original_md5_hash):
+        return
+
+    logger.warning("File corrupted. Re-downloading...")
+    os.remove(check_target)
+
+    if retry_on_mismatch:
+        _download_verified_asset(
+            filename=filename,
+            original_md5_hash=original_md5_hash,
+            destination=destination,
+            check_target=check_target,
+            retry_on_mismatch=False,
+        )
+        return
+
+    raise ValueError(f"Downloaded asset {filename!r} failed MD5 verification.")
+
+
+def download_assets(
+    asset_name: Assets | str,
+    directory: str | Path | None = None,
+) -> str:
+    """
+    Download a specified asset if it doesn't already exist or is corrupted.
+
+    Args:
+        asset_name: The name or type of the asset to be downloaded.
+        directory: Optional output directory. Defaults to the current working
+            directory for backward compatibility.
+
+    Returns:
+        The downloaded asset path. When `directory` is omitted, this preserves
+        the historical filename-only return value.
+
+    Example:
+        ```pycon
+        >>> from supervision.assets import download_assets, ImageAssets, VideoAssets
+        >>> download_assets(VideoAssets.VEHICLES)  # doctest: +SKIP
+        'vehicles.mp4'
+
+        >>> download_assets(ImageAssets.PEOPLE_WALKING)  # doctest: +SKIP
+        'people-walking.jpg'
+
+        ```
+    """
+
+    filename = asset_name.filename if isinstance(asset_name, Assets) else asset_name
+    if directory is None:
+        destination = Path.cwd() / filename
+        check_target: str | Path = filename
+        return_value = filename
+    else:
+        destination_directory = Path(directory).expanduser().resolve()
+        destination = destination_directory / filename
+        check_target = str(destination)
+        return_value = str(destination)
+
+    if filename in MEDIA_ASSETS:
+        original_md5_hash = MEDIA_ASSETS[filename][1]
+        if not Path(check_target).exists():
+            logger.info("Downloading %s assets", filename)
+            _download_verified_asset(
+                filename=filename,
+                original_md5_hash=original_md5_hash,
+                destination=destination,
+                check_target=check_target,
+            )
+        else:
+            if not is_md5_hash_matching(check_target, original_md5_hash):
+                logger.warning("File corrupted. Re-downloading...")
+                os.remove(check_target)
+                _download_verified_asset(
+                    filename=filename,
+                    original_md5_hash=original_md5_hash,
+                    destination=destination,
+                    check_target=check_target,
+                )
+
+            logger.info("%s asset download complete.", filename)
+    else:
+        valid_assets = ", ".join(filename for filename in MEDIA_ASSETS.keys())
+        raise ValueError(
+            f"Invalid asset. It should be one of the following: {valid_assets}."
+        )
+
+    return return_value