lightly-studio 0.3.1__py3-none-any.whl

lightly_studio/metrics/detection/map.py

@@ -0,0 +1,268 @@
+"""Implementation for calculating Mean Average Precision (MAP)."""
+
+from __future__ import annotations
+
+import math
+from collections import defaultdict
+from collections.abc import Sequence
+from uuid import UUID
+
+import torch
+from pydantic import BaseModel
+from torch import Tensor
+from torchmetrics.detection.mean_ap import MeanAveragePrecision
+
+from lightly_studio.models.annotation.annotation_base import AnnotationBaseTable
+
+
+class DetectionMetricsMAP(BaseModel):
+    """Response for computing the MAP detection metric."""
+
+    name: str
+    map: float
+    map_small: float
+    map_medium: float
+    map_large: float
+    mar_1: float
+    mar_10: float
+    mar_100: float
+    mar_small: float
+    mar_medium: float
+    mar_large: float
+    map_50: float
+    map_75: float
+    map_per_class: dict[str, float] | None = None
+    mar_100_per_class: dict[str, float] | None = None
+    classes: list[int]
+
+    @property
+    def value(self) -> float:
+        """Backwards compatibility alias for map."""
+        return self.map
+
+    @property
+    def per_class(self) -> dict[str, float] | None:
+        """Backwards compatibility alias for map_per_class."""
+        return self.map_per_class
+
+
+def calculate_map_metric(  # noqa: C901
+    pred_annotations: Sequence[AnnotationBaseTable],
+    gt_annotations: Sequence[AnnotationBaseTable],
+) -> DetectionMetricsMAP:
+    """Calculate the Mean Average Precision (MAP) metric."""
+    if not gt_annotations or not pred_annotations:
+        return DetectionMetricsMAP(
+            name="mAP@.5:.95",
+            map=0.0,
+            map_small=0.0,
+            map_medium=0.0,
+            map_large=0.0,
+            mar_1=0.0,
+            mar_10=0.0,
+            mar_100=0.0,
+            mar_small=0.0,
+            mar_medium=0.0,
+            mar_large=0.0,
+            map_50=-1.0,
+            map_75=-1.0,
+            map_per_class=None,
+            mar_100_per_class=None,
+            classes=[],
+        )
+    uuid_to_int_label: dict[UUID, int] = {}
+
+    # Map sample_id to annotations.
+    sample_id_to_annotations = _group_by_sample_id(
+        pred_annotations=pred_annotations,
+        gt_annotations=gt_annotations,
+    )
+
+    # Create the MeanAveragePrecision object.
+    map_metric = MeanAveragePrecision(
+        box_format="xywh",
+        iou_thresholds=None,
+        class_metrics=True,
+        average="macro",
+        backend="faster_coco_eval",
+    )
+
+    # Compute MAP updating the object sample by sample.
+    for _sample_id, (
+        sample_pred_annotations,
+        sample_gt_annotations,
+    ) in sample_id_to_annotations.items():
+        # Convert to torchmetrics format.
+        prediction_tm = _convert_to_torchmetrics(
+            annotations=sample_pred_annotations,
+            is_prediction=True,
+            uuid_to_int_label_map=uuid_to_int_label,
+        )
+        ground_truth_tm = _convert_to_torchmetrics(
+            annotations=sample_gt_annotations,
+            is_prediction=False,
+            uuid_to_int_label_map=uuid_to_int_label,
+        )
+
+        # Update the metric.
+        map_metric.update(preds=[prediction_tm], target=[ground_truth_tm])
+
+    # Compute the final results.
+    results = map_metric.compute()
+    # Invert label map to recover original UUIDs
+    int_to_uuid = {v: k for k, v in uuid_to_int_label.items()}
+
+    # Helper to convert Tensor to Python list
+    def to_list(t: Tensor) -> list[float]:
+        if t.dim() == 0:
+            return [t.item()]
+        return t.tolist()
+
+    # Scalar metrics
+    map_val = results["map"].item()
+    map_small = results["map_small"].item()
+    map_medium = results["map_medium"].item()
+    map_large = results["map_large"].item()
+    mar_1 = results["mar_1"].item()
+    mar_10 = results["mar_10"].item()
+    mar_100 = results["mar_100"].item()
+    mar_small = results["mar_small"].item()
+    mar_medium = results["mar_medium"].item()
+    mar_large = results["mar_large"].item()
+
+    # IoU-specific metrics
+    map_50 = results.get("map_50")
+    map_50 = map_50.item() if isinstance(map_50, Tensor) and map_50.dim() == 0 else -1.0
+    map_75 = results.get("map_75")
+    map_75 = map_75.item() if isinstance(map_75, Tensor) and map_75.dim() == 0 else -1.0
+
+    # Per-class metrics
+    raw_map_pc = results.get("map_per_class")
+    per_class_dict: dict[str, float] | None = None
+    if raw_map_pc is not None:
+        pc_list = to_list(raw_map_pc)
+        per_class_dict = {}
+        for idx, val in enumerate(pc_list):
+            if not math.isnan(val):
+                per_class_dict[str(int_to_uuid[idx])] = val
+
+    # Per-class recall at 100 detections
+    raw_mar100_pc = results.get("mar_100_per_class")
+    mar_100_pc_dict: dict[str, float] | None = None
+    if raw_mar100_pc is not None:
+        m100_list = to_list(raw_mar100_pc)
+        mar_100_pc_dict = {}
+        for idx, val in enumerate(m100_list):
+            if not math.isnan(val):
+                mar_100_pc_dict[str(int_to_uuid[idx])] = val
+
+    # Observed classes
+    raw_classes = results.get("classes")
+    classes_list_output: list[int] = []
+    if raw_classes is not None:
+        classes_list_output = [int(x) for x in to_list(raw_classes)]
+
+    return DetectionMetricsMAP(
+        name="mAP@.5:.95",
+        map=map_val,
+        map_small=map_small,
+        map_medium=map_medium,
+        map_large=map_large,
+        mar_1=mar_1,
+        mar_10=mar_10,
+        mar_100=mar_100,
+        mar_small=mar_small,
+        mar_medium=mar_medium,
+        mar_large=mar_large,
+        map_50=map_50,
+        map_75=map_75,
+        map_per_class=per_class_dict,
+        mar_100_per_class=mar_100_pc_dict,
+        classes=classes_list_output,
+    )
+
+
+def _group_by_sample_id(
+    pred_annotations: Sequence[AnnotationBaseTable],
+    gt_annotations: Sequence[AnnotationBaseTable],
+) -> dict[
+    UUID,
+    tuple[list[AnnotationBaseTable], list[AnnotationBaseTable]],
+]:
+    """Group prediction and ground truth annotations by sample_id.
+
+    Returns a dictionary with sample_id as key and a tuple of
+    (list of prediction annotations, list of ground truth annotations) as value.
+    """
+    sample_id_to_annotations: defaultdict[
+        UUID,
+        tuple[list[AnnotationBaseTable], list[AnnotationBaseTable]],
+    ] = defaultdict(lambda: ([], []))
+    for ann in pred_annotations:
+        sample_id_to_annotations[ann.sample_id][0].append(ann)
+    for ann in gt_annotations:
+        sample_id_to_annotations[ann.sample_id][1].append(ann)
+    return sample_id_to_annotations
+
+
+def _convert_to_torchmetrics(
+    annotations: Sequence[AnnotationBaseTable],
+    is_prediction: bool,
+    uuid_to_int_label_map: dict[UUID, int],
+) -> dict[str, Tensor]:
+    """Convert annotations to torchmetrics format.
+
+    Args:
+        annotations: List of bounding box annotations.
+        is_prediction: Whether the annotations are predictions.
+        uuid_to_int_label_map: Map from UUID to integer label.
+
+    Returns:
+        Dictionary in torchmetrics format. For predictions the keys are
+        `boxes`, `scores`, `labels`. For ground truth they are `boxes`,
+        `labels`.
+
+    """
+    if not annotations:
+        empty_boxes = torch.empty((0, 4), dtype=torch.float32)
+        empty_labels = torch.empty((0,), dtype=torch.int64)
+        if is_prediction:
+            return {
+                "boxes": empty_boxes,
+                "scores": torch.empty((0,), dtype=torch.float32),
+                "labels": empty_labels,
+            }
+        return {"boxes": empty_boxes, "labels": empty_labels}
+
+    boxes = torch.tensor(
+        [
+            [
+                a.object_detection_details.x,
+                a.object_detection_details.y,
+                a.object_detection_details.width,
+                a.object_detection_details.height,
+            ]
+            for a in annotations
+            if a.object_detection_details is not None
+        ],
+        dtype=torch.float32,
+    )
+
+    mapped_labels = []
+    # Use the passed-in map instead of a global one
+    for uuid_label in (a.annotation_label_id for a in annotations):
+        if uuid_label not in uuid_to_int_label_map:
+            # Assign the next available integer ID based on the current map size
+            uuid_to_int_label_map[uuid_label] = len(uuid_to_int_label_map)
+        mapped_labels.append(uuid_to_int_label_map[uuid_label])
+
+    labels = torch.tensor(mapped_labels, dtype=torch.int64)
+
+    if is_prediction:
+        scores = torch.tensor([a.confidence for a in annotations], dtype=torch.float32)
+        return {
+            "boxes": boxes,
+            "scores": scores,
+            "labels": labels,
+        }
+    return {"boxes": boxes, "labels": labels}

lightly_studio/models/__init__.py

@@ -0,0 +1 @@
+

lightly_studio/models/annotation/annotation_base.py

@@ -0,0 +1,171 @@
+"""This module defines the base annotation model."""
+
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, List, Optional
+from uuid import UUID, uuid4
+
+from pydantic import BaseModel
+from pydantic import Field as PydanticField
+from sqlalchemy.orm import Mapped
+from sqlmodel import Field, Relationship, SQLModel
+
+from lightly_studio.models.annotation.instance_segmentation import (
+    InstanceSegmentationAnnotationTable,
+    InstanceSegmentationAnnotationView,
+)
+from lightly_studio.models.annotation.links import AnnotationTagLinkTable
+from lightly_studio.models.annotation.object_detection import (
+    ObjectDetectionAnnotationTable,
+    ObjectDetectionAnnotationView,
+)
+from lightly_studio.models.annotation.semantic_segmentation import (
+    SemanticSegmentationAnnotationTable,
+    SemanticSegmentationAnnotationView,
+)
+from lightly_studio.models.annotation_task import (
+    AnnotationTaskTable,
+    AnnotationType,
+)
+
+if TYPE_CHECKING:
+    from lightly_studio.models.annotation_label import (
+        AnnotationLabelTable,
+    )
+    from lightly_studio.models.sample import (
+        SampleTable,
+    )
+    from lightly_studio.models.tag import TagTable
+else:
+    TagTable = object
+    SampleTable = object
+    AnnotationLabelTable = object
+
+
+class AnnotationBaseTable(SQLModel, table=True):
+    """Base class for all annotation models."""
+
+    __tablename__ = "annotation_base"
+
+    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc), index=True)
+
+    annotation_id: UUID = Field(default_factory=uuid4, primary_key=True)
+    annotation_type: AnnotationType
+    annotation_label_id: UUID = Field(foreign_key="annotation_labels.annotation_label_id")
+    annotation_task_id: UUID = Field(
+        foreign_key="annotation_tasks.annotation_task_id",
+    )
+    confidence: Optional[float] = None
+    dataset_id: UUID = Field(foreign_key="datasets.dataset_id")
+    sample_id: UUID = Field(foreign_key="samples.sample_id")
+
+    annotation_label: Mapped["AnnotationLabelTable"] = Relationship(
+        sa_relationship_kwargs={"lazy": "select"},
+    )
+    annotation_task: Mapped["AnnotationTaskTable"] = Relationship(
+        sa_relationship_kwargs={"lazy": "select"},
+    )
+    sample: Mapped[Optional["SampleTable"]] = Relationship(
+        sa_relationship_kwargs={"lazy": "select"},
+    )
+    tags: Mapped[List["TagTable"]] = Relationship(
+        back_populates="annotations",
+        link_model=AnnotationTagLinkTable,
+    )
+
+    """ Details about object detection. """
+    object_detection_details: Mapped[Optional["ObjectDetectionAnnotationTable"]] = Relationship(
+        back_populates="annotation_base",
+        sa_relationship_kwargs={"lazy": "select"},
+    )
+
+    """ Details about instance segmentation. """
+    instance_segmentation_details: Mapped[Optional["InstanceSegmentationAnnotationTable"]] = (
+        Relationship(
+            back_populates="annotation_base",
+            sa_relationship_kwargs={"lazy": "select"},
+        )
+    )
+
+    """ Details about semantic segmentation. """
+    semantic_segmentation_details: Mapped[Optional["SemanticSegmentationAnnotationTable"]] = (
+        Relationship(
+            back_populates="annotation_base",
+            sa_relationship_kwargs={"lazy": "select"},
+        )
+    )
+
+
+class AnnotationCreate(SQLModel):
+    """Input model for creating annotations."""
+
+    """ Required properties for all annotations. """
+    annotation_label_id: UUID
+    annotation_type: AnnotationType
+    annotation_task_id: UUID
+    confidence: Optional[float] = None
+    dataset_id: UUID
+    sample_id: UUID
+
+    """ Optional properties for object detection. """
+    x: Optional[float] = None
+    y: Optional[float] = None
+    width: Optional[float] = None
+    height: Optional[float] = None
+
+    """ Optional properties for instance and semantic segmentation. """
+    segmentation_mask: Optional[List[int]] = None
+
+
+class AnnotationSampleView(SQLModel):
+    """Sample class for annotation view."""
+
+    file_path_abs: str
+    file_name: str
+    dataset_id: UUID
+    sample_id: UUID
+    width: int
+    height: int
+
+
+class AnnotationView(SQLModel):
+    """Response model for bounding box annotation."""
+
+    class AnnotationLabel(SQLModel):
+        """Model used when retrieving an annotation label."""
+
+        annotation_label_name: str
+
+    sample_id: UUID
+    dataset_id: UUID
+    annotation_id: UUID
+    annotation_type: AnnotationType
+    annotation_label: AnnotationLabel
+    annotation_task_id: UUID
+    confidence: Optional[float] = None
+
+    object_detection_details: Optional[ObjectDetectionAnnotationView] = None
+    instance_segmentation_details: Optional[InstanceSegmentationAnnotationView] = None
+    semantic_segmentation_details: Optional[SemanticSegmentationAnnotationView] = None
+
+
+class AnnotationWithSampleView(AnnotationView):
+    """Response model for bounding box annotation."""
+
+    sample: AnnotationSampleView
+
+
+class AnnotationViewsWithCount(BaseModel):
+    """Response model for counted annotations."""
+
+    annotations: List[AnnotationWithSampleView] = PydanticField(..., alias="data")
+    total_count: int
+    next_cursor: Optional[int] = PydanticField(..., alias="nextCursor")
+
+    class Config:  # noqa: D106
+        populate_by_name = True
+
+
+class AnnotationDetailsView(AnnotationView):
+    """Representing detailed view of an annotation."""
+
+    sample: AnnotationSampleView

lightly_studio/models/annotation/instance_segmentation.py

@@ -0,0 +1,56 @@
+"""Instance segmentation annotation models.
+
+Instance segmentation combines object detection and semantic segmentation,
+identifying objects and providing pixel-level masks for each instance.
+"""
+
+from typing import TYPE_CHECKING, List, Optional
+from uuid import UUID, uuid4
+
+from sqlalchemy import ARRAY, Column, Integer
+from sqlalchemy.orm import Mapped
+from sqlmodel import Field, Relationship, SQLModel
+
+if TYPE_CHECKING:
+    from lightly_studio.models.annotation.annotation_base import (
+        AnnotationBaseTable,
+    )
+else:
+    AnnotationBaseTable = object
+
+
+class InstanceSegmentationAnnotationTable(SQLModel, table=True):
+    """Database table model for instance segmentation annotations."""
+
+    __tablename__ = "instance_segmentation_annotations"
+
+    annotation_id: UUID = Field(
+        default_factory=uuid4,
+        primary_key=True,
+        foreign_key="annotation_base.annotation_id",
+    )
+
+    annotation_base: Mapped["AnnotationBaseTable"] = Relationship(
+        back_populates="instance_segmentation_details"
+    )
+
+    x: float
+    y: float
+    width: float
+    height: float
+    # TODO(Kondrat 06/2025): We need to fix logic in the loader,
+    # because it shouldn't be optional.
+    # lightly_studio/dataset/loader.py#L148
+    segmentation_mask: Optional[List[int]] = Field(
+        default=None, sa_column=Column(ARRAY(Integer), nullable=True)
+    )
+
+
+class InstanceSegmentationAnnotationView(SQLModel):
+    """API response model for instance segmentation annotations."""
+
+    x: float
+    y: float
+    width: float
+    height: float
+    segmentation_mask: Optional[List[int]] = None

lightly_studio/models/annotation/links.py

@@ -0,0 +1,17 @@
+"""This module defines the base annotation model."""
+
+from typing import Optional
+from uuid import UUID
+
+from sqlmodel import Field, SQLModel
+
+
+class AnnotationTagLinkTable(SQLModel, table=True):
+    """Model defines the link table between annotations and tags."""
+
+    annotation_id: Optional[UUID] = Field(
+        default=None,
+        foreign_key="annotation_base.annotation_id",
+        primary_key=True,
+    )
+    tag_id: Optional[UUID] = Field(default=None, foreign_key="tags.tag_id", primary_key=True)

lightly_studio/models/annotation/object_detection.py

@@ -0,0 +1,47 @@
+"""Object detection annotation models.
+
+Object detection identifies and locates objects in images using bounding boxes.
+"""
+
+from typing import TYPE_CHECKING
+from uuid import UUID, uuid4
+
+from sqlalchemy.orm import Mapped
+from sqlmodel import Field, Relationship, SQLModel
+
+if TYPE_CHECKING:
+    from lightly_studio.models.annotation.annotation_base import (
+        AnnotationBaseTable,
+    )
+else:
+    AnnotationBaseTable = object
+
+
+class ObjectDetectionAnnotationTable(SQLModel, table=True):
+    """Database table model for object detection annotations."""
+
+    __tablename__ = "object_detection_annotations"
+
+    annotation_id: UUID = Field(
+        default_factory=uuid4,
+        primary_key=True,
+        foreign_key="annotation_base.annotation_id",
+    )
+
+    annotation_base: Mapped["AnnotationBaseTable"] = Relationship(
+        back_populates="object_detection_details"
+    )
+
+    x: float
+    y: float
+    width: float
+    height: float
+
+
+class ObjectDetectionAnnotationView(SQLModel):
+    """API response model for object detection annotations."""
+
+    x: float
+    y: float
+    width: float
+    height: float

lightly_studio/models/annotation/semantic_segmentation.py

@@ -0,0 +1,44 @@
+"""This module defines the semantic segmentation annotation model.
+
+Semantic segmentation is a computer vision task that assigns a class label to
+each pixel in an image. This module provides the data models for storing and
+managing semantic segmentation annotations.
+"""
+
+from typing import TYPE_CHECKING, List
+from uuid import UUID, uuid4
+
+from sqlalchemy import ARRAY, Column, Integer
+from sqlalchemy.orm import Mapped
+from sqlmodel import Field, Relationship, SQLModel
+
+if TYPE_CHECKING:
+    from lightly_studio.models.annotation.annotation_base import (
+        AnnotationBaseTable,
+    )
+else:
+    AnnotationBaseTable = object
+
+
+class SemanticSegmentationAnnotationTable(SQLModel, table=True):
+    """Model used to define semantic segmentation annotation table."""
+
+    __tablename__ = "semantic_segmentation_annotations"
+
+    annotation_id: UUID = Field(
+        default_factory=uuid4,
+        primary_key=True,
+        foreign_key="annotation_base.annotation_id",
+    )
+
+    segmentation_mask: List[int] = Field(sa_column=Column(ARRAY(Integer), nullable=True))
+
+    annotation_base: Mapped["AnnotationBaseTable"] = Relationship(
+        back_populates="semantic_segmentation_details"
+    )
+
+
+class SemanticSegmentationAnnotationView(SQLModel):
+    """Response model for semantic segmentation annotation."""
+
+    segmentation_mask: List[int]

lightly_studio/models/annotation_label.py

@@ -0,0 +1,47 @@
+"""This module defines the AnnotationLabel model for the application."""
+
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, List
+from uuid import UUID, uuid4
+
+from sqlalchemy.orm import Mapped
+from sqlmodel import Field, Relationship, SQLModel
+
+if TYPE_CHECKING:
+    from lightly_studio.models.annotation.annotation_base import (
+        AnnotationBaseTable,
+    )
+
+else:
+    AnnotationBaseTable = object
+
+
+class AnnotationLabelBase(SQLModel):
+    """Base class for the AnnotationLabel model."""
+
+    annotation_label_name: str
+
+
+class AnnotationLabelCreate(AnnotationLabelBase):
+    """Model used when creating an annotation label."""
+
+
+class AnnotationLabelView(AnnotationLabelBase):
+    """Model used when retrieving an annotation label."""
+
+    annotation_label_id: UUID
+
+
+class AnnotationLabelTable(AnnotationLabelBase, table=True):
+    """This class defines the AnnotationLabel model."""
+
+    __tablename__ = "annotation_labels"
+
+    annotation_label_id: UUID = Field(default_factory=uuid4, primary_key=True)
+    created_at: str = Field(
+        default_factory=lambda: datetime.now(timezone.utc),
+        index=True,
+    )
+    annotations: Mapped[List["AnnotationBaseTable"]] = Relationship(
+        back_populates="annotation_label",
+    )

lightly_studio/models/annotation_task.py

@@ -0,0 +1,28 @@
+"""This module defines the AnnotationTask model."""
+
+from datetime import datetime
+from enum import Enum
+from uuid import UUID, uuid4
+
+from sqlmodel import Field, SQLModel
+
+
+class AnnotationType(str, Enum):
+    """The type of annotation task."""
+
+    BBOX = "bbox"
+    CLASSIFICATION = "classification"
+    SEMANTIC_SEGMENTATION = "semantic_segmentation"
+    INSTANCE_SEGMENTATION = "instance_segmentation"
+    OBJECT_DETECTION = "object_detection"
+
+
+class AnnotationTaskTable(SQLModel, table=True):
+    """The annotation task model."""
+
+    __tablename__ = "annotation_tasks"
+    annotation_task_id: UUID = Field(default_factory=uuid4, primary_key=True)
+    name: str
+    created_at: datetime = Field(default_factory=datetime.utcnow)
+    annotation_type: AnnotationType
+    is_prediction: bool