lightly-studio 0.4.6__py3-none-any.whl

lightly_studio/metadata/gps_coordinate.py

@@ -0,0 +1,41 @@
+"""GPS coordinate representation for complex metadata."""
+
+from typing import Dict
+
+
+class GPSCoordinate:
+    """Represents a GPS coordinate."""
+
+    def __init__(self, lat: float, lon: float):
+        """Initialize GPS coordinate.
+
+        Args:
+            lat: Latitude in decimal degrees.
+            lon: Longitude in decimal degrees.
+        """
+        self.lat = lat
+        self.lon = lon
+
+    def __repr__(self) -> str:
+        """String representation of the GPS coordinate."""
+        return f"GPSCoordinate(lat={self.lat}, lon={self.lon})"
+
+    def as_dict(self) -> Dict[str, float]:
+        """Convert the GPSCoordinate to a dictionary.
+
+        Returns:
+            Dictionary with 'lat' and 'lon' keys.
+        """
+        return {"lat": self.lat, "lon": self.lon}
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, float]) -> "GPSCoordinate":
+        """Create a GPSCoordinate from a dictionary.
+
+        Args:
+            data: Dictionary with 'lat' and 'lon' keys.
+
+        Returns:
+            A GPSCoordinate instance.
+        """
+        return cls(lat=data["lat"], lon=data["lon"])

lightly_studio/metadata/metadata_protocol.py

@@ -0,0 +1,17 @@
+"""Protocol for complex metadata types that can be stored in JSON columns."""
+
+from typing import Any, Dict, Protocol, runtime_checkable
+
+
+@runtime_checkable
+class ComplexMetadata(Protocol):
+    """Protocol for complex types that can be serialized to/from JSON."""
+
+    def as_dict(self) -> Dict[str, Any]:
+        """Convert the complex metadata to a dictionary for JSON storage."""
+        ...
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ComplexMetadata":
+        """Create the complex metadata from a dictionary."""
+        ...

lightly_studio/models/__init__.py

@@ -0,0 +1 @@
+

lightly_studio/models/annotation/annotation_base.py

@@ -0,0 +1,303 @@
+"""This module defines the base annotation model."""
+
+from datetime import datetime, timezone
+from enum import Enum
+from typing import TYPE_CHECKING, List, Optional, Union
+from uuid import UUID
+
+from pydantic import BaseModel, ConfigDict
+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.dataset import SampleType
+from lightly_studio.models.sample import SampleTable
+from lightly_studio.models.video import VideoFrameTable
+
+if TYPE_CHECKING:
+    from lightly_studio.models.annotation_label import (
+        AnnotationLabelTable,
+    )
+    from lightly_studio.models.image import ImageTable
+    from lightly_studio.models.tag import TagTable
+
+else:
+    TagTable = object
+    AnnotationLabelTable = object
+    ImageTable = object
+
+
+class AnnotationType(str, Enum):
+    """The type of annotation task."""
+
+    CLASSIFICATION = "classification"
+    SEMANTIC_SEGMENTATION = "semantic_segmentation"
+    INSTANCE_SEGMENTATION = "instance_segmentation"
+    OBJECT_DETECTION = "object_detection"
+
+
+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)
+
+    sample_id: UUID = Field(foreign_key="sample.sample_id", primary_key=True)
+    annotation_type: AnnotationType
+    annotation_label_id: UUID = Field(foreign_key="annotation_label.annotation_label_id")
+
+    confidence: Optional[float] = None
+    parent_sample_id: UUID = Field(foreign_key="sample.sample_id")
+
+    annotation_label: Mapped["AnnotationLabelTable"] = Relationship(
+        sa_relationship_kwargs={"lazy": "select"},
+    )
+    sample: Mapped["SampleTable"] = Relationship(
+        sa_relationship_kwargs={
+            "lazy": "select",
+            "foreign_keys": "[AnnotationBaseTable.sample_id]",
+        },
+    )
+    parent_sample: Mapped[Optional["SampleTable"]] = Relationship(
+        back_populates="annotations",
+        sa_relationship_kwargs={
+            "lazy": "select",
+            "foreign_keys": "[AnnotationBaseTable.parent_sample_id]",
+        },
+    )
+    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
+    confidence: Optional[float] = None
+    parent_sample_id: UUID
+
+    """ Optional properties for object detection. """
+    x: Optional[int] = None
+    y: Optional[int] = None
+    width: Optional[int] = None
+    height: Optional[int] = None
+
+    """ Optional properties for instance and semantic segmentation. """
+    segmentation_mask: Optional[List[int]] = None
+
+
+class AnnotationView(BaseModel):
+    """Response model for bounding box annotation."""
+
+    model_config = ConfigDict(populate_by_name=True, from_attributes=True)
+
+    class AnnotationLabel(SQLModel):
+        """Model used when retrieving an annotation label."""
+
+        annotation_label_name: str
+
+    class AnnotationViewTag(SQLModel):
+        """Tag view inside Annotation view."""
+
+        tag_id: UUID
+        name: str
+
+    parent_sample_id: UUID
+    sample_id: UUID
+    annotation_type: AnnotationType
+    annotation_label: AnnotationLabel
+    confidence: Optional[float] = None
+    created_at: datetime
+
+    object_detection_details: Optional[ObjectDetectionAnnotationView] = None
+    instance_segmentation_details: Optional[InstanceSegmentationAnnotationView] = None
+    semantic_segmentation_details: Optional[SemanticSegmentationAnnotationView] = None
+
+    tags: List[AnnotationViewTag] = []
+
+
+class AnnotationViewsWithCount(BaseModel):
+    """Response model for counted annotations."""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    annotations: List[AnnotationView] = PydanticField(..., alias="data")
+    total_count: int
+    next_cursor: Optional[int] = PydanticField(..., alias="nextCursor")
+
+
+class SampleAnnotationView(BaseModel):
+    """Response model for sample annotation view."""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    dataset_id: UUID
+
+
+class ImageAnnotationView(BaseModel):
+    """Response model for image annotation view."""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    sample_id: UUID
+    file_path_abs: str
+    width: int
+    height: int
+    sample: SampleAnnotationView
+
+
+class VideoAnnotationView(BaseModel):
+    """Response model for video view."""
+
+    height: int
+    width: int
+    file_path_abs: str
+
+
+class VideoFrameAnnotationView(BaseModel):
+    """Response model for video frame annotation view."""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    sample_id: UUID
+    video: VideoAnnotationView
+
+
+class AnnotationWithPayloadView(BaseModel):
+    """Response model for annotation with payload."""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    parent_sample_type: SampleType
+    annotation: AnnotationView
+    parent_sample_data: Union[ImageAnnotationView, VideoFrameAnnotationView]
+
+
+class AnnotationWithPayloadAndCountView(BaseModel):
+    """Response model for counted annotations with payload."""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    annotations: List[AnnotationWithPayloadView] = PydanticField(..., alias="data")
+    total_count: int
+    next_cursor: Optional[int] = PydanticField(None, alias="nextCursor")
+
+
+class SampleAnnotationDetailsView(BaseModel):
+    """Response model for sample annotation details view."""
+
+    sample_id: UUID
+    dataset_id: UUID
+    tags: List["TagTable"] = []
+
+    @classmethod
+    def from_sample_table(cls, sample: SampleTable) -> "SampleAnnotationDetailsView":
+        """Convert sample table to sample annotation details view."""
+        return SampleAnnotationDetailsView(
+            sample_id=sample.sample_id,
+            tags=sample.tags,
+            dataset_id=sample.dataset_id,
+        )
+
+
+class ImageAnnotationDetailsView(BaseModel):
+    """Response model for image annotation details view."""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    file_path_abs: str
+    file_name: str
+    width: int
+    height: int
+    sample: SampleAnnotationDetailsView
+
+    @classmethod
+    def from_image_table(cls, image: ImageTable) -> "ImageAnnotationDetailsView":
+        """Convert image table to image annotation details view."""
+        return ImageAnnotationDetailsView(
+            height=image.height,
+            width=image.width,
+            file_path_abs=image.file_path_abs,
+            file_name=image.file_name,
+            sample=SampleAnnotationDetailsView.from_sample_table(image.sample),
+        )
+
+
+class VideoFrameAnnotationDetailsView(BaseModel):
+    """Response model for video frame annotation view."""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    sample_id: UUID
+    frame_number: int
+    frame_timestamp_s: float
+
+    video: VideoAnnotationView
+    sample: SampleAnnotationDetailsView
+
+    @classmethod
+    def from_video_frame_table(
+        cls, video_frame: VideoFrameTable
+    ) -> "VideoFrameAnnotationDetailsView":
+        """Convert video frame table to video frame annotation details view."""
+        return VideoFrameAnnotationDetailsView(
+            sample_id=video_frame.sample_id,
+            frame_number=video_frame.frame_number,
+            frame_timestamp_s=video_frame.frame_timestamp_s,
+            video=VideoAnnotationView(
+                width=video_frame.video.width,
+                height=video_frame.video.height,
+                file_path_abs=video_frame.video.file_path_abs,
+            ),
+            sample=SampleAnnotationDetailsView.from_sample_table(video_frame.sample),
+        )
+
+
+class AnnotationDetailsWithPayloadView(BaseModel):
+    """Response model for annotation details with payload."""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    parent_sample_type: SampleType
+    annotation: AnnotationView
+    parent_sample_data: Union[ImageAnnotationDetailsView, VideoFrameAnnotationDetailsView]

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_annotation"
+
+    sample_id: UUID = Field(
+        default_factory=uuid4,
+        primary_key=True,
+        foreign_key="annotation_base.sample_id",
+    )
+
+    annotation_base: Mapped["AnnotationBaseTable"] = Relationship(
+        back_populates="instance_segmentation_details"
+    )
+
+    x: int
+    y: int
+    width: int
+    height: int
+    # 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: int
+    y: int
+    width: int
+    height: int
+    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_sample_id: Optional[UUID] = Field(
+        default=None,
+        foreign_key="annotation_base.sample_id",
+        primary_key=True,
+    )
+    tag_id: Optional[UUID] = Field(default=None, foreign_key="tag.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_annotation"
+
+    sample_id: UUID = Field(
+        default_factory=uuid4,
+        primary_key=True,
+        foreign_key="annotation_base.sample_id",
+    )
+
+    annotation_base: Mapped["AnnotationBaseTable"] = Relationship(
+        back_populates="object_detection_details"
+    )
+
+    x: int
+    y: int
+    width: int
+    height: int
+
+
+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_annotation"
+
+    sample_id: UUID = Field(
+        default_factory=uuid4,
+        primary_key=True,
+        foreign_key="annotation_base.sample_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 = Field(unique=True)
+
+
+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_label"
+
+    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/caption.py

@@ -0,0 +1,49 @@
+"""This module defines the caption model."""
+
+from datetime import datetime, timezone
+from uuid import UUID
+
+from sqlalchemy.orm import Mapped
+from sqlmodel import Field, Relationship, SQLModel
+
+from lightly_studio.models.sample import SampleTable
+
+
+class CaptionTable(SQLModel, table=True):
+    """Class for caption model."""
+
+    __tablename__ = "caption"
+
+    sample_id: UUID = Field(foreign_key="sample.sample_id", primary_key=True)
+    parent_sample_id: UUID = Field(foreign_key="sample.sample_id")
+    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc), index=True)
+    text: str
+
+    sample: Mapped["SampleTable"] = Relationship(
+        sa_relationship_kwargs={
+            "lazy": "select",
+            "foreign_keys": "[CaptionTable.sample_id]",
+        },
+    )
+    parent_sample: Mapped["SampleTable"] = Relationship(
+        back_populates="captions",
+        sa_relationship_kwargs={
+            "lazy": "select",
+            "foreign_keys": "[CaptionTable.parent_sample_id]",
+        },
+    )
+
+
+class CaptionCreate(SQLModel):
+    """Input model for creating captions."""
+
+    parent_sample_id: UUID
+    text: str
+
+
+class CaptionView(SQLModel):
+    """Response model for caption."""
+
+    parent_sample_id: UUID
+    sample_id: UUID
+    text: str

lightly_studio/models/classifier.py

@@ -0,0 +1,20 @@
+"""This module defines the data model for the classifier."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from pydantic import BaseModel
+
+
+class EmbeddingClassifier(BaseModel):
+    """Base class for the Classifier model."""
+
+    """The name of the classifier."""
+    classifier_name: str
+
+    """The ID of the classifier."""
+    classifier_id: UUID
+
+    """List of classes supported by the classifier."""
+    class_list: list[str]