lightly-studio 0.4.6__py3-none-any.whl

lightly_studio/models/tag.py

@@ -0,0 +1,96 @@
+"""This module contains the Tag model and related enumerations."""
+
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, List, Literal, Optional
+from uuid import UUID, uuid4
+
+from sqlalchemy import UniqueConstraint
+from sqlalchemy.orm import Mapped
+from sqlmodel import Field, Relationship, SQLModel, String
+
+from lightly_studio.models.annotation.links import AnnotationTagLinkTable
+from lightly_studio.models.sample import SampleTagLinkTable
+
+if TYPE_CHECKING:
+    from lightly_studio.models.annotation.annotation_base import (
+        AnnotationBaseTable,
+    )
+    from lightly_studio.models.sample import SampleTable
+
+else:
+    SampleTable = object
+    TagTable = object
+    AnnotationBaseTable = object
+
+
+# TagKind is the kind of tag we support.
+TagKind = Literal[
+    "sample",
+    "annotation",
+]
+
+
+class TagBase(SQLModel):
+    """Base class for the Tag model."""
+
+    name: str
+    description: Optional[str] = ""
+    kind: TagKind = "sample"
+
+
+class TagCreate(TagBase):
+    """Tag model when creating."""
+
+    dataset_id: UUID
+
+
+class TagCreateBody(TagBase):
+    """Tag model when creating."""
+
+
+class TagUpdate(TagBase):
+    """Tag model when updating."""
+
+
+class TagUpdateBody(TagBase):
+    """Tag model when updating."""
+
+    dataset_id: Optional[UUID] = None
+
+
+class TagView(TagBase):
+    """Tag model when retrieving."""
+
+    tag_id: UUID
+    kind: TagKind
+    created_at: datetime
+    updated_at: datetime
+
+
+class TagTable(TagBase, table=True):
+    """This class defines the Tag model."""
+
+    __tablename__ = "tag"
+    # ensure there can only be one tag named "lightly_studio" per dataset
+    __table_args__ = (
+        UniqueConstraint("dataset_id", "kind", "name", name="unique_name_constraint"),
+    )
+    tag_id: UUID = Field(default_factory=uuid4, primary_key=True)
+    dataset_id: UUID
+    kind: TagKind = Field(sa_type=String)
+    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc), index=True)
+    updated_at: datetime = Field(
+        default_factory=lambda: datetime.now(timezone.utc),
+    )
+
+    """The sample ids associated with the tag."""
+    samples: Mapped[List["SampleTable"]] = Relationship(
+        back_populates="tags",
+        link_model=SampleTagLinkTable,
+    )
+
+    """The annotation ids associated with the tag (legacy bounding box)."""
+    annotations: Mapped[List["AnnotationBaseTable"]] = Relationship(
+        back_populates="tags",
+        link_model=AnnotationTagLinkTable,
+    )

lightly_studio/models/two_dim_embedding.py

@@ -0,0 +1,16 @@
+"""Database table storing cached 2D embeddings."""
+
+from __future__ import annotations
+
+from sqlalchemy import ARRAY, Float
+from sqlmodel import Column, Field, SQLModel
+
+
+class TwoDimEmbeddingTable(SQLModel, table=True):
+    """Persisted 2D embedding projection identified by a deterministic hash."""
+
+    __tablename__ = "two_dim_embeddings"
+
+    hash: str = Field(primary_key=True)
+    x: list[float] = Field(sa_column=Column(ARRAY(Float)))
+    y: list[float] = Field(sa_column=Column(ARRAY(Float)))

lightly_studio/models/video.py

@@ -0,0 +1,161 @@
+"""This module defines the Video and VideoFrame model for the application."""
+
+from typing import TYPE_CHECKING, List, Optional
+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.range import FloatRange, IntRange
+
+if TYPE_CHECKING:
+    from lightly_studio.models.annotation.annotation_base import (
+        AnnotationBaseTable,
+    )
+    from lightly_studio.models.sample import SampleTable, SampleView
+
+else:
+    AnnotationBaseTable = object
+    SampleTable = object
+    SampleView = object
+
+
+class VideoBase(SQLModel):
+    """Base class for the Video model."""
+
+    """The width of the video in pixels."""
+    width: int
+
+    """The height of the video in pixels."""
+    height: int
+
+    """The duration of the video in seconds."""
+    duration_s: Optional[float] = Field(default=None)
+
+    """The fps of the video."""
+    fps: float
+
+    """The file name of the video."""
+    file_name: str
+
+    """The path of the video."""
+    file_path_abs: str = Field(default=None)
+
+
+class VideoCreate(VideoBase):
+    """Video class when inserting."""
+
+
+class VideoTable(VideoBase, table=True):
+    """This class defines the Video ORM table."""
+
+    __tablename__ = "video"
+    sample_id: UUID = Field(foreign_key="sample.sample_id", primary_key=True)
+    frames: Mapped[List["VideoFrameTable"]] = Relationship(back_populates="video")
+    sample: Mapped["SampleTable"] = Relationship()
+
+
+class VideoView(SQLModel):
+    """Video class when retrieving."""
+
+    width: int
+    height: int
+    duration_s: Optional[float] = None
+    fps: float
+    file_name: str
+    file_path_abs: str
+    sample_id: UUID
+    sample: SampleView
+    frame: Optional["FrameView"] = None
+
+
+class VideoViewsWithCount(BaseModel):
+    """Response model for counted videos."""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    samples: List[VideoView] = PydanticField(..., alias="data")
+    total_count: int
+    next_cursor: Optional[int] = PydanticField(None, alias="nextCursor")
+
+
+class VideoFrameBase(SQLModel):
+    """Base class for the VideoFrame model."""
+
+    """The frame number of the video frame."""
+    frame_number: int
+
+    """Presentation timestamp in the video's time_base and in seconds."""
+    frame_timestamp_pts: int
+    frame_timestamp_s: float
+
+    """The video ID to which the video frame belongs."""
+    parent_sample_id: UUID = Field(default=None, foreign_key="video.sample_id")
+
+    """The rotation of the encoded frame in degrees.
+
+    Valid values are 0, 90, 180, 270. This field is needed for video access with OpenCV
+    which does not read rotation metadata correctly."""
+    rotation_deg: int = 0
+
+
+class VideoFrameCreate(VideoFrameBase):
+    """VideoFrame class when inserting."""
+
+
+class VideoFrameTable(VideoFrameBase, table=True):
+    """This class defines the VideoFrame model."""
+
+    __tablename__ = "video_frame"
+    sample_id: UUID = Field(foreign_key="sample.sample_id", primary_key=True)
+
+    sample: Mapped["SampleTable"] = Relationship()
+    video: Mapped["VideoTable"] = Relationship(back_populates="frames")
+
+
+class VideoFrameView(SQLModel):
+    """VideoFrame class when retrieving."""
+
+    frame_number: int
+    frame_timestamp_s: float
+    sample_id: UUID
+
+    # Video metadata routed from parent video
+    video: VideoView
+    sample: "SampleView"
+
+
+class FrameView(SQLModel):
+    """VideoFrame class when retrieving."""
+
+    frame_number: int
+    frame_timestamp_s: float
+    sample_id: UUID
+    sample: SampleView
+
+
+class VideoFrameViewsWithCount(BaseModel):
+    """Response model for counted video frames."""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    samples: List[VideoFrameView] = PydanticField(..., alias="data")
+    total_count: int
+    next_cursor: Optional[int] = PydanticField(None, alias="nextCursor")
+
+
+class VideoFieldsBoundsView(BaseModel):
+    """Response model for the video fields bounds."""
+
+    width: IntRange
+    height: IntRange
+    duration_s: FloatRange
+    fps: FloatRange
+
+
+class VideoFrameFieldsBoundsView(BaseModel):
+    """Response model for the video frame fields bounds."""
+
+    frame_number: IntRange

lightly_studio/plugins/base_operator.py

@@ -0,0 +1,60 @@
+"""Base operator class for LightlyStudio plugins."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any
+from uuid import UUID
+
+from sqlmodel import Session
+
+from lightly_studio.plugins.parameter import BaseParameter
+
+
+@dataclass
+class OperatorResult:
+    """Result returned by operator execution."""
+
+    success: bool
+    message: str
+
+
+class BaseOperator(ABC):
+    """Base class for all operators."""
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """Return the operator name."""
+
+    @property
+    @abstractmethod
+    def description(self) -> str:
+        """Return the description of the operator."""
+
+    @property
+    @abstractmethod
+    def parameters(self) -> list[BaseParameter]:
+        """Return the list of parameters this operator expects."""
+
+    @abstractmethod
+    def execute(
+        self,
+        *,
+        session: Session,
+        dataset_id: UUID,
+        parameters: dict[str, Any],
+    ) -> OperatorResult:
+        """Execute the operator with the given parameters.
+
+        Args:
+            session: Database session.
+            dataset_id: ID of the dataset to operate on.
+            parameters: Parameters passed to the operator.
+
+        Returns:
+            Dictionary with 'success' (bool) and 'message' (str) keys.
+        """
+        # TODO (Jonas 11/2025): The parameters dict should be validated against self.parameters,
+        # for now we leave it to the operator implementation.

lightly_studio/plugins/operator_registry.py

@@ -0,0 +1,47 @@
+"""Operator registry for LightlyStudio plugins."""
+
+from __future__ import annotations
+
+import uuid
+from dataclasses import dataclass
+
+from .base_operator import BaseOperator
+
+
+@dataclass
+class RegisteredOperatorMetadata:
+    """Meta data for a registered operator."""
+
+    operator_id: str
+    name: str
+
+
+class OperatorRegistry:
+    """Registry for managing operators."""
+
+    def __init__(self) -> None:
+        """Initialize the operator registry."""
+        self._operators: dict[str, BaseOperator] = {}
+
+    def register(self, operator: BaseOperator) -> None:
+        """Register an operator."""
+        operator_id = str(uuid.uuid4())
+        self._operators[operator_id] = operator
+
+    def get_all_metadata(self) -> list[RegisteredOperatorMetadata]:
+        """Get all registered operators with their names."""
+        return [
+            RegisteredOperatorMetadata(
+                operator_id=operator_id,
+                name=operator.name,
+            )
+            for operator_id, operator in self._operators.items()
+        ]
+
+    def get_by_id(self, operator_id: str) -> BaseOperator | None:
+        """Get an operator by its ID."""
+        return self._operators.get(operator_id)
+
+
+# Global registry instance
+operator_registry = OperatorRegistry()

lightly_studio/plugins/parameter.py

@@ -0,0 +1,70 @@
+"""Parameter for operators for LightlyStudio plugins."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, Generic, TypeVar, cast
+
+T = TypeVar("T")
+
+
+@dataclass
+class BaseParameter(ABC):
+    """Base parameter definition shared across operator parameters."""
+
+    name: str
+    description: str = ""
+    default: Any = None
+    required: bool = True
+    param_type: str | None = None
+
+    def __post_init__(self) -> None:
+        """Run value validation once the dataclass is initialized."""
+        if self.default is not None:
+            self.default = self._validate(self.default)
+
+    @abstractmethod
+    def _validate(self, value: Any) -> Any:
+        """Validate the parameter value."""
+
+
+class BuiltinParameter(BaseParameter, Generic[T]):
+    """Represents a built-in operator parameter."""
+
+    def __post_init__(self) -> None:
+        """Set up type information and validate default value."""
+        if not hasattr(self, "_parameter_type") or self._parameter_type is None:
+            raise NotImplementedError("Subclasses must define _parameter_type class attribute")
+        self._type = self._parameter_type
+        self.param_type = self._parameter_type.__name__
+        super().__post_init__()
+
+    def _validate(self, value: T) -> T:
+        if isinstance(value, self._type):
+            return cast(T, value)
+        raise TypeError(f"Expected value of type '{self._type.__name__}' but got {type(value)}'")
+
+
+class IntParameter(BuiltinParameter[int]):
+    """Represents an integer operator parameter."""
+
+    _parameter_type = int
+
+
+class FloatParameter(BuiltinParameter[float]):
+    """Represents a float operator parameter."""
+
+    _parameter_type = float
+
+
+class BoolParameter(BuiltinParameter[bool]):
+    """Represents a boolean operator parameter."""
+
+    _parameter_type = bool
+
+
+class StringParameter(BuiltinParameter[str]):
+    """Represents a string operator parameter."""
+
+    _parameter_type = str

lightly_studio/resolvers/annotation_label_resolver/__init__.py

@@ -0,0 +1,22 @@
+"""Handler for database operations related to annotation labels."""
+
+from .create import create
+from .delete import delete
+from .get_all import get_all, get_all_sorted_alphabetically
+from .get_by_id import get_by_id
+from .get_by_ids import get_by_ids
+from .get_by_label_name import get_by_label_name
+from .names_by_ids import names_by_ids
+from .update import update
+
+__all__ = [
+    "create",
+    "delete",
+    "get_all",
+    "get_all_sorted_alphabetically",
+    "get_by_id",
+    "get_by_ids",
+    "get_by_label_name",
+    "names_by_ids",
+    "update",
+]

lightly_studio/resolvers/annotation_label_resolver/create.py

@@ -0,0 +1,27 @@
+"""Create annotation label functionality."""
+
+from __future__ import annotations
+
+from sqlmodel import Session
+
+from lightly_studio.models.annotation_label import (
+    AnnotationLabelCreate,
+    AnnotationLabelTable,
+)
+
+
+def create(session: Session, label: AnnotationLabelCreate) -> AnnotationLabelTable:
+    """Create a new annotation label in the database.
+
+    Args:
+        session (Session): The database session.
+        label (AnnotationLabelCreate): The annotation label data to be created.
+
+    Returns:
+        AnnotationLabelTable: The newly created annotation label record.
+    """
+    db_label = AnnotationLabelTable.model_validate(label)
+    session.add(db_label)
+    session.commit()
+    session.refresh(db_label)
+    return db_label

lightly_studio/resolvers/annotation_label_resolver/delete.py

@@ -0,0 +1,28 @@
+"""Delete annotation label functionality."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlmodel import Session
+
+from .get_by_id import get_by_id
+
+
+def delete(session: Session, label_id: UUID) -> bool:
+    """Delete an annotation label.
+
+    Args:
+        session (Session): The database session.
+        label_id (UUID): The unique identifier of the annotation label.
+
+    Returns:
+        bool: True if the label was deleted, False if the label was not found.
+    """
+    label = get_by_id(session=session, label_id=label_id)
+    if not label:
+        return False
+
+    session.delete(label)
+    session.commit()
+    return True

lightly_studio/resolvers/annotation_label_resolver/get_all.py

@@ -0,0 +1,37 @@
+"""Get all annotation labels functionality."""
+
+from __future__ import annotations
+
+from sqlmodel import Session, col, select
+
+from lightly_studio.models.annotation_label import AnnotationLabelTable
+
+
+def get_all(session: Session) -> list[AnnotationLabelTable]:
+    """Retrieve all annotation labels.
+
+    Args:
+        session (Session): The database session.
+
+    Returns:
+        list[AnnotationLabelTable]: A list of annotation labels.
+    """
+    labels = session.exec(
+        select(AnnotationLabelTable).order_by(col(AnnotationLabelTable.created_at).asc())
+    ).all()
+    return list(labels) if labels else []
+
+
+def get_all_sorted_alphabetically(session: Session) -> list[AnnotationLabelTable]:
+    """Retrieve all annotation labels sorted alphabetically.
+
+    Args:
+        session (Session): The database session.
+
+    Returns:
+        list[AnnotationLabelTable]: A list of annotation labels.
+    """
+    labels = session.exec(
+        select(AnnotationLabelTable).order_by(col(AnnotationLabelTable.annotation_label_name).asc())
+    ).all()
+    return list(labels) if labels else []

lightly_studio/resolvers/annotation_label_resolver/get_by_id.py

@@ -0,0 +1,24 @@
+"""Get annotation label by ID functionality."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlmodel import Session, select
+
+from lightly_studio.models.annotation_label import AnnotationLabelTable
+
+
+def get_by_id(session: Session, label_id: UUID) -> AnnotationLabelTable | None:
+    """Retrieve a single annotation label by ID.
+
+    Args:
+        session (Session): The database session used to execute the query.
+        label_id (UUID): The unique identifier of the annotation label to get.
+
+    Returns:
+        AnnotationLabelTable | None: The annotation label if found, or None.
+    """
+    return session.exec(
+        select(AnnotationLabelTable).where(AnnotationLabelTable.annotation_label_id == label_id)
+    ).one_or_none()

lightly_studio/resolvers/annotation_label_resolver/get_by_ids.py

@@ -0,0 +1,25 @@
+"""Get annotation labels by IDs functionality."""
+
+from __future__ import annotations
+
+from typing import Sequence
+from uuid import UUID
+
+from sqlmodel import Session, col, select
+
+from lightly_studio.models.annotation_label import AnnotationLabelTable
+
+
+def get_by_ids(session: Session, ids: Sequence[UUID]) -> list[AnnotationLabelTable]:
+    """Retrieve annotation labels by their IDs.
+
+    Output order matches the input order.
+    """
+    results = session.exec(
+        select(AnnotationLabelTable).where(
+            col(AnnotationLabelTable.annotation_label_id).in_(list(ids))
+        )
+    ).all()
+    # Return labels in the same order as the input ids.
+    label_map = {label.annotation_label_id: label for label in results}
+    return [label_map[id_] for id_ in ids if id_ in label_map]

lightly_studio/resolvers/annotation_label_resolver/get_by_label_name.py

@@ -0,0 +1,24 @@
+"""Handler for database operations related to annotation labels."""
+
+from __future__ import annotations
+
+from sqlmodel import Session, select
+
+from lightly_studio.models.annotation_label import (
+    AnnotationLabelTable,
+)
+
+
+def get_by_label_name(session: Session, label_name: str) -> AnnotationLabelTable | None:
+    """Retrieve a single annotation label by its name.
+
+    Args:
+        session: The database session to use for the query.
+        label_name: The name of the annotation label to retrieve.
+
+    Returns:
+        The AnnotationLabelTable instance if found, None otherwise.
+    """
+    return session.exec(
+        select(AnnotationLabelTable).where(AnnotationLabelTable.annotation_label_name == label_name)
+    ).one_or_none()

lightly_studio/resolvers/annotation_label_resolver/names_by_ids.py

@@ -0,0 +1,25 @@
+"""Get annotation label names by IDs functionality."""
+
+from __future__ import annotations
+
+from typing import Sequence
+from uuid import UUID
+
+from sqlmodel import Session
+
+from .get_by_ids import get_by_ids
+
+
+def names_by_ids(session: Session, ids: Sequence[UUID]) -> dict[str, str]:
+    """Return a dictionary mapping annotation label IDs to their names.
+
+    Args:
+        session (Session): The database session used to query the labels.
+        ids (Sequence[UUID]): A sequence of UUIDs of the annotation label IDs.
+
+    Returns:
+        dict[str, str]: A dictionary with string representations of the UUIDs
+        and the values are the corresponding annotation label names.
+    """
+    labels = get_by_ids(session=session, ids=ids)
+    return {str(label.annotation_label_id): label.annotation_label_name for label in labels}