valor-lite 0.37.1__py3-none-any.whl

valor_lite/object_detection/annotation.py

@@ -0,0 +1,238 @@
+from dataclasses import dataclass, field
+from typing import Any, Generic, TypeVar
+
+import numpy as np
+from numpy.typing import NDArray
+from shapely.geometry import Polygon as ShapelyPolygon
+
+
+@dataclass
+class BoundingBox:
+    """
+    Represents a bounding box with associated labels and optional scores.
+
+    Parameters
+    ----------
+    uid : str
+        A unique identifier.
+    xmin : float
+        The minimum x-coordinate of the bounding box.
+    xmax : float
+        The maximum x-coordinate of the bounding box.
+    ymin : float
+        The minimum y-coordinate of the bounding box.
+    ymax : float
+        The maximum y-coordinate of the bounding box.
+    labels : list of str
+        List of labels associated with the bounding box.
+    scores : list of float, optional
+        Confidence scores corresponding to each label. Defaults to an empty list.
+    metadata : dict[str, Any], optional
+        A dictionary containing any metadata to be used within filtering operations.
+
+    Examples
+    --------
+    Ground Truth Example:
+
+    >>> bbox = BoundingBox(uid="xyz", xmin=10.0, xmax=50.0, ymin=20.0, ymax=60.0, labels=['cat'])
+
+    Prediction Example:
+
+    >>> bbox = BoundingBox(
+    ...     uid="abc",
+    ...     xmin=10.0, xmax=50.0, ymin=20.0, ymax=60.0,
+    ...     labels=['cat', 'dog'], scores=[0.9, 0.1]
+    ... )
+    """
+
+    uid: str
+    xmin: float
+    xmax: float
+    ymin: float
+    ymax: float
+    labels: list[str]
+    scores: list[float] = field(default_factory=list)
+    metadata: dict[str, Any] | None = None
+
+    def __post_init__(self):
+        if len(self.scores) == 0 and len(self.labels) != 1:
+            raise ValueError(
+                "Ground truths must be defined with no scores and a single label. If you meant to define a prediction, then please include one score for every label provided."
+            )
+        if len(self.scores) > 0 and len(self.labels) != len(self.scores):
+            raise ValueError(
+                "If scores are defined, there must be a 1:1 pairing with labels."
+            )
+
+    @property
+    def extrema(self) -> tuple[float, float, float, float]:
+        """
+        Returns the annotation's data representation.
+
+        Returns
+        -------
+        tuple[float, float, float, float]
+            A tuple in the form (xmin, xmax, ymin, ymax).
+        """
+        return (self.xmin, self.xmax, self.ymin, self.ymax)
+
+
+@dataclass
+class Polygon:
+    """
+    Represents a polygon shape with associated labels and optional scores.
+
+    Parameters
+    ----------
+    uid : str
+        A unique identifier.
+    shape : shapely.geometry.Polygon
+        A Shapely polygon object representing the shape.
+    labels : list of str
+        List of labels associated with the polygon.
+    scores : list of float, optional
+        Confidence scores corresponding to each label. Defaults to an empty list.
+    metadata : dict[str, Any], optional
+        A dictionary containing any metadata to be used within filtering operations.
+
+    Examples
+    --------
+    Ground Truth Example:
+
+    >>> from shapely.geometry import Polygon as ShapelyPolygon
+    >>> shape = ShapelyPolygon([(0, 0), (1, 0), (1, 1), (0, 1)])
+    >>> polygon = Polygon(uid="xyz", shape=shape, labels=['building'])
+
+    Prediction Example:
+
+    >>> polygon = Polygon(
+    ...     uid="abc", shape=shape, labels=['building'], scores=[0.95]
+    ... )
+    """
+
+    uid: str
+    shape: ShapelyPolygon
+    labels: list[str]
+    scores: list[float] = field(default_factory=list)
+    metadata: dict[str, Any] | None = None
+
+    def __post_init__(self):
+        if not isinstance(self.shape, ShapelyPolygon):
+            raise TypeError("shape must be of type shapely.geometry.Polygon.")
+        if self.shape.is_empty:
+            raise ValueError("Polygon is empty.")
+
+        if len(self.scores) == 0 and len(self.labels) != 1:
+            raise ValueError(
+                "Ground truths must be defined with no scores and a single label. If you meant to define a prediction, then please include one score for every label provided."
+            )
+        if len(self.scores) > 0 and len(self.labels) != len(self.scores):
+            raise ValueError(
+                "If scores are defined, there must be a 1:1 pairing with labels."
+            )
+
+
+@dataclass
+class Bitmask:
+    """
+    Represents a binary mask with associated labels and optional scores.
+
+    Parameters
+    ----------
+    uid : str
+        A unique identifier.
+    mask : NDArray[np.bool_]
+        A NumPy array of boolean values representing the mask.
+    labels : list of str
+        List of labels associated with the mask.
+    scores : list of float, optional
+        Confidence scores corresponding to each label. Defaults to an empty list.
+    metadata : dict[str, Any], optional
+        A dictionary containing any metadata to be used within filtering operations.
+
+    Examples
+    --------
+    Ground Truth Example:
+
+    >>> import numpy as np
+    >>> mask = np.array([[True, False], [False, True]], dtype=np.bool_)
+    >>> bitmask = Bitmask(uid="abc", mask=mask, labels=['tree'])
+
+    Prediction Example:
+
+    >>> bitmask = Bitmask(
+    ...     uid="xyz", mask=mask, labels=['tree'], scores=[0.85]
+    ... )
+    """
+
+    uid: str
+    mask: NDArray[np.bool_]
+    labels: list[str]
+    scores: list[float] = field(default_factory=list)
+    metadata: dict[str, Any] | None = None
+
+    def __post_init__(self):
+
+        if (
+            not isinstance(self.mask, np.ndarray)
+            or self.mask.dtype != np.bool_
+        ):
+            raise ValueError(
+                "Expected mask to be of type `NDArray[np.bool_]`."
+            )
+        elif not self.mask.any():
+            raise ValueError("Mask does not define any object instances.")
+
+        if len(self.scores) == 0 and len(self.labels) != 1:
+            raise ValueError(
+                "Ground truths must be defined with no scores and a single label. If you meant to define a prediction, then please include one score for every label provided."
+            )
+        if len(self.scores) > 0 and len(self.labels) != len(self.scores):
+            raise ValueError(
+                "If scores are defined, there must be a 1:1 pairing with labels."
+            )
+
+
+AnnotationType = TypeVar("AnnotationType", BoundingBox, Polygon, Bitmask)
+
+
+@dataclass
+class Detection(Generic[AnnotationType]):
+    """
+    Detection data structure holding ground truths and predictions for object detection tasks.
+
+    Parameters
+    ----------
+    uid : str
+        Unique identifier for the image or sample.
+    groundtruths : list[BoundingBox] | list[Polygon] | list[Bitmask]
+        List of ground truth annotations.
+    predictions : list[BoundingBox] | list[Polygon] | list[Bitmask]
+        List of predicted annotations.
+    metadata : dict[str, Any], optional
+        A dictionary containing any metadata to be used within filtering operations.
+
+    Examples
+    --------
+    >>> bbox_gt = BoundingBox(xmin=10, xmax=50, ymin=20, ymax=60, labels=['cat'])
+    >>> bbox_pred = BoundingBox(
+    ...     xmin=12, xmax=48, ymin=22, ymax=58, labels=['cat'], scores=[0.9]
+    ... )
+    >>> detection = Detection(
+    ...     uid='image_001',
+    ...     groundtruths=[bbox_gt],
+    ...     predictions=[bbox_pred]
+    ... )
+    """
+
+    uid: str
+    groundtruths: list[AnnotationType]
+    predictions: list[AnnotationType]
+    metadata: dict[str, Any] | None = None
+
+    def __post_init__(self):
+        for prediction in self.predictions:
+            if len(prediction.scores) != len(prediction.labels):
+                raise ValueError(
+                    "Predictions must provide a score for every label."
+                )