lightly-studio 0.4.6__py3-none-any.whl

lightly_studio/resolvers/tag_resolver.py

@@ -0,0 +1,299 @@
+"""Handler for database operations related to tags."""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from uuid import UUID
+
+import sqlmodel
+from sqlmodel import Session, col, select
+
+from lightly_studio.models.annotation.annotation_base import AnnotationBaseTable
+from lightly_studio.models.annotation.links import AnnotationTagLinkTable
+from lightly_studio.models.sample import SampleTable, SampleTagLinkTable
+from lightly_studio.models.tag import TagCreate, TagTable, TagUpdate
+
+
+def create(session: Session, tag: TagCreate) -> TagTable:
+    """Create a new tag in the database."""
+    db_tag = TagTable.model_validate(tag)
+    session.add(db_tag)
+    session.commit()
+    session.refresh(db_tag)
+    return db_tag
+
+
+# TODO(Michal, 06/2025): Use Paginated struct instead of offset/limit.
+def get_all_by_dataset_id(
+    session: Session, dataset_id: UUID, offset: int = 0, limit: int | None = None
+) -> list[TagTable]:
+    """Retrieve all tags with pagination."""
+    query = (
+        select(TagTable)
+        .where(TagTable.dataset_id == dataset_id)
+        .order_by(col(TagTable.created_at).asc(), col(TagTable.tag_id).asc())
+        .offset(offset)
+    )
+    if limit is not None:
+        query = query.limit(limit)
+    tags = session.exec(query).all()
+    return list(tags) if tags else []
+
+
+def get_by_id(session: Session, tag_id: UUID) -> TagTable | None:
+    """Retrieve a single tag by ID."""
+    return session.exec(select(TagTable).where(TagTable.tag_id == tag_id)).one_or_none()
+
+
+def get_by_name(session: Session, tag_name: str, dataset_id: UUID | None) -> TagTable | None:
+    """Retrieve a single tag by ID."""
+    if dataset_id:
+        return session.exec(
+            select(TagTable)
+            .where(TagTable.dataset_id == dataset_id)
+            .where(TagTable.name == tag_name)
+        ).one_or_none()
+    return session.exec(select(TagTable).where(TagTable.name == tag_name)).one_or_none()
+
+
+def update(session: Session, tag_id: UUID, tag_data: TagUpdate) -> TagTable | None:
+    """Update an existing tag."""
+    tag = get_by_id(session=session, tag_id=tag_id)
+    if not tag:
+        return None
+
+    # due to duckdb/OLAP optimisations, update operations effecting unique
+    # constraints (e.g colums) will lead to a unique constraint violation.
+    # This is due to a update is implemented as delete+insert. The error
+    # happens only within the same session.
+    # To fix it, we can delete, commit + insert a new tag.
+    # https://duckdb.org/docs/sql/indexes#over-eager-unique-constraint-checking
+    session.delete(tag)
+    session.commit()
+
+    # create clone of tag with updated values
+    tag_updated = TagTable.model_validate(tag)
+    tag_updated.name = tag_data.name
+    tag_updated.description = tag_data.description
+    tag_updated.updated_at = datetime.now(timezone.utc)
+
+    session.add(tag_updated)
+    session.commit()
+    session.refresh(tag_updated)
+    return tag_updated
+
+
+def delete(session: Session, tag_id: UUID) -> bool:
+    """Delete a tag."""
+    tag = get_by_id(session=session, tag_id=tag_id)
+    if not tag:
+        return False
+
+    session.delete(tag)
+    session.commit()
+    return True
+
+
+def add_tag_to_sample(
+    session: Session,
+    tag_id: UUID,
+    sample: SampleTable,
+) -> SampleTable | None:
+    """Add a tag to a sample."""
+    tag = get_by_id(session=session, tag_id=tag_id)
+    if not tag or not tag.tag_id:
+        return None
+    if tag.kind != "sample":
+        raise ValueError(f"Tag {tag_id} is not of kind 'sample'")
+
+    sample.tags.append(tag)
+    session.add(sample)
+    session.commit()
+    session.refresh(sample)
+    return sample
+
+
+def remove_tag_from_sample(
+    session: Session,
+    tag_id: UUID,
+    sample: SampleTable,
+) -> SampleTable | None:
+    """Remove a tag from a sample."""
+    tag = get_by_id(session=session, tag_id=tag_id)
+    if not tag or not tag.tag_id:
+        return None
+    if tag.kind != "sample":
+        raise ValueError(f"Tag {tag_id} is not of kind 'sample'")
+
+    sample.tags.remove(tag)
+    session.add(sample)
+    session.commit()
+    session.refresh(sample)
+    return sample
+
+
+def add_tag_to_annotation(
+    session: Session,
+    tag_id: UUID,
+    annotation: AnnotationBaseTable,
+) -> AnnotationBaseTable | None:
+    """Add a tag to a annotation."""
+    tag = get_by_id(session=session, tag_id=tag_id)
+    if not tag or not tag.tag_id:
+        return None
+    if tag.kind != "annotation":
+        raise ValueError(f"Tag {tag_id} is not of kind 'annotation'")
+
+    annotation.tags.append(tag)
+    session.add(annotation)
+    session.commit()
+    session.refresh(annotation)
+    return annotation
+
+
+def assign_tag_to_annotation(
+    session: Session,
+    tag: TagTable,
+    annotation: AnnotationBaseTable,
+) -> AnnotationBaseTable:
+    """Add a tag to a annotation."""
+    annotation.tags.append(tag)
+    session.add(annotation)
+    session.commit()
+    session.refresh(annotation)
+    return annotation
+
+
+def remove_tag_from_annotation(
+    session: Session,
+    tag_id: UUID,
+    annotation: AnnotationBaseTable,
+) -> AnnotationBaseTable | None:
+    """Remove a tag from a annotation."""
+    tag = get_by_id(session=session, tag_id=tag_id)
+    if not tag or not tag.tag_id:
+        return None
+    if tag.kind != "annotation":
+        raise ValueError(f"Tag {tag_id} is not of kind 'annotation'")
+
+    annotation.tags.remove(tag)
+    session.add(annotation)
+    session.commit()
+    session.refresh(annotation)
+    return annotation
+
+
+def add_sample_ids_to_tag_id(
+    session: Session,
+    tag_id: UUID,
+    sample_ids: list[UUID],
+) -> TagTable | None:
+    """Add a list of sample_ids to a tag."""
+    tag = get_by_id(session=session, tag_id=tag_id)
+    if not tag or not tag.tag_id:
+        return None
+    if tag.kind != "sample":
+        raise ValueError(f"Tag {tag_id} is not of kind 'sample'")
+
+    for sample_id in sample_ids:
+        session.merge(SampleTagLinkTable(sample_id=sample_id, tag_id=tag_id))
+
+    session.commit()
+    session.refresh(tag)
+    return tag
+
+
+def remove_sample_ids_from_tag_id(
+    session: Session,
+    tag_id: UUID,
+    sample_ids: list[UUID],
+) -> TagTable | None:
+    """Remove a list of sample_ids to a tag."""
+    tag = get_by_id(session=session, tag_id=tag_id)
+    if not tag or not tag.tag_id:
+        return None
+    if tag.kind != "sample":
+        raise ValueError(f"Tag {tag_id} is not of kind 'sample'")
+
+    session.exec(  # type:ignore[call-overload]
+        sqlmodel.delete(SampleTagLinkTable).where(
+            col(SampleTagLinkTable.tag_id) == tag_id,
+            col(SampleTagLinkTable.sample_id).in_(sample_ids),
+        )
+    )
+
+    session.commit()
+    session.refresh(tag)
+    return tag
+
+
+def add_annotation_ids_to_tag_id(
+    session: Session,
+    tag_id: UUID,
+    annotation_ids: list[UUID],
+) -> TagTable | None:
+    """Add a list of annotation_ids to a tag."""
+    tag = get_by_id(session=session, tag_id=tag_id)
+    if not tag or not tag.tag_id:
+        return None
+    if tag.kind != "annotation":
+        raise ValueError(f"Tag {tag_id} is not of kind 'annotation'")
+
+    for annotation_id in annotation_ids:
+        session.merge(
+            AnnotationTagLinkTable(
+                tag_id=tag_id,
+                annotation_sample_id=annotation_id,
+            )
+        )
+
+    session.commit()
+    session.refresh(tag)
+    return tag
+
+
+def remove_annotation_ids_from_tag_id(
+    session: Session,
+    tag_id: UUID,
+    annotation_ids: list[UUID],
+) -> TagTable | None:
+    """Remove a list of things to a tag."""
+    tag = get_by_id(session=session, tag_id=tag_id)
+    if not tag or not tag.tag_id:
+        return None
+    if tag.kind != "annotation":
+        raise ValueError(f"Tag {tag_id} is not of kind 'annotation'")
+
+    session.exec(  # type:ignore[call-overload]
+        sqlmodel.delete(AnnotationTagLinkTable).where(
+            col(AnnotationTagLinkTable.tag_id) == tag_id,
+            col(AnnotationTagLinkTable.annotation_sample_id).in_(annotation_ids),
+        )
+    )
+
+    session.commit()
+    session.refresh(tag)
+    return tag
+
+
+def get_or_create_sample_tag_by_name(
+    session: Session,
+    dataset_id: UUID,
+    tag_name: str,
+) -> TagTable:
+    """Get an existing sample tag by name or create a new one if it doesn't exist.
+
+    Args:
+        session: Database session for executing queries.
+        dataset_id: The dataset ID to search/create the tag for.
+        tag_name: Name of the tag to get or create.
+
+    Returns:
+        The existing or newly created sample tag.
+    """
+    existing_tag = get_by_name(session=session, tag_name=tag_name, dataset_id=dataset_id)
+    if existing_tag:
+        return existing_tag
+
+    new_tag = TagCreate(name=tag_name, dataset_id=dataset_id, kind="sample")
+    return create(session=session, tag=new_tag)

lightly_studio/resolvers/twodim_embedding_resolver.py

@@ -0,0 +1,119 @@
+"""Handler for getting cached 2D embeddings from high-dimensional embeddings."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+import numpy as np
+from lightly_mundig import TwoDimEmbedding  # type: ignore[import-untyped]
+from numpy.typing import NDArray
+from sqlmodel import Session, col, select
+
+from lightly_studio.dataset.env import LIGHTLY_STUDIO_LICENSE_KEY
+from lightly_studio.models.embedding_model import EmbeddingModelTable
+from lightly_studio.models.sample import SampleTable
+from lightly_studio.models.two_dim_embedding import TwoDimEmbeddingTable
+from lightly_studio.resolvers import sample_embedding_resolver
+
+
+def get_twodim_embeddings(
+    session: Session,
+    dataset_id: UUID,
+    embedding_model_id: UUID,
+) -> tuple[NDArray[np.float32], NDArray[np.float32], list[UUID]]:
+    """Return cached 2D embeddings together with their sample identifiers.
+
+    Uses a cache to avoid recomputing the 2D embeddings. The cache key combines the sorted
+    sample identifiers with a deterministic 64-bit hash over the stored high-dimensional
+    embeddings.
+
+    Args:
+        session: Database session.
+        dataset_id: Dataset identifier.
+        embedding_model_id: Embedding model identifier.
+
+    Returns:
+        Tuple of (x coordinates, y coordinates, ordered sample IDs).
+    """
+    embedding_model = session.get(EmbeddingModelTable, embedding_model_id)
+    if embedding_model is None:
+        raise ValueError(f"Embedding model {embedding_model_id} not found.")
+
+    # Define a fixed order of sample IDs for the cache key.
+    sample_ids_ordered = list(
+        session.exec(
+            select(SampleTable.sample_id)
+            .where(SampleTable.dataset_id == dataset_id)
+            .order_by(col(SampleTable.sample_id).asc())
+        ).all()
+    )
+
+    # Check if we have a cached 2D embedding for the given samples and embedding model.
+    # The order is defined by sample_ids_ordered.
+    cache_key, sample_ids_of_samples_with_embeddings = (
+        sample_embedding_resolver.get_hash_by_sample_ids(
+            session=session,
+            sample_ids_ordered=sample_ids_ordered,
+            embedding_model_id=embedding_model_id,
+        )
+    )
+
+    if not sample_ids_of_samples_with_embeddings:
+        empty = np.array([], dtype=np.float32)
+        return empty, empty, []
+
+    # If there is a cached entry, return it.
+    cached = session.get(TwoDimEmbeddingTable, cache_key)
+    if cached is not None:
+        x_values = np.array(cached.x, dtype=np.float32)
+        y_values = np.array(cached.y, dtype=np.float32)
+        return x_values, y_values, sample_ids_of_samples_with_embeddings
+
+    # No cached entry found - load the high-dimensional embeddings.
+    # The order is defined by sample_ids_of_samples_with_embeddings.
+    sample_embeddings = sample_embedding_resolver.get_by_sample_ids(
+        session=session,
+        sample_ids=sample_ids_of_samples_with_embeddings,
+        embedding_model_id=embedding_model_id,
+    )
+
+    # If there are no embeddings, return empty arrays.
+    if not sample_embeddings:
+        empty = np.array([], dtype=np.float32)
+        return empty, empty, []
+
+    # Compute the 2D embedding from the high-dimensional embeddings.
+    # The order is now defined by sample_embeddings. They are the ordered subset of the
+    # sample_ids_of_samples_with_embeddings that have embeddings.
+    sample_ids_of_samples_with_embeddings = [embedding.sample_id for embedding in sample_embeddings]
+    embedding_values = [embedding.embedding for embedding in sample_embeddings]
+    planar_embeddings = _calculate_2d_embeddings(embedding_values)
+    embeddings_2d = np.asarray(planar_embeddings, dtype=np.float32)
+    x_values, y_values = embeddings_2d[:, 0], embeddings_2d[:, 1]
+
+    # Write the computed 2D embeddings to the cache.
+    cache_entry = TwoDimEmbeddingTable(hash=cache_key, x=list(x_values), y=list(y_values))
+    session.add(cache_entry)
+    session.commit()
+
+    return x_values, y_values, sample_ids_of_samples_with_embeddings
+
+
+def _calculate_2d_embeddings(embedding_values: list[list[float]]) -> list[tuple[float, float]]:
+    n_samples = len(embedding_values)
+    # For 0, 1 or 2 samples we hard-code deterministic coordinates.
+    if n_samples == 0:
+        return []
+    if n_samples == 1:
+        return [(0.0, 0.0)]
+    if n_samples == 2:  # noqa: PLR2004
+        return [(0.0, 0.0), (1.0, 1.0)]
+
+    license_key = LIGHTLY_STUDIO_LICENSE_KEY
+    if license_key is None:
+        raise ValueError(
+            "LIGHTLY_STUDIO_LICENSE_KEY environment variable is not set. "
+            "Please set it to your LightlyStudio license key."
+        )
+    embedding_calculator = TwoDimEmbedding(embedding_values, license_key)
+    return embedding_calculator.calculate_2d_embedding()  # type: ignore[no-any-return]

lightly_studio/resolvers/video_frame_resolver/__init__.py

@@ -0,0 +1,23 @@
+"""Resolvers for video_frame database operations."""
+
+from lightly_studio.resolvers.video_frame_resolver.count_video_frames_annotations import (
+    count_video_frames_annotations,
+)
+from lightly_studio.resolvers.video_frame_resolver.create_many import create_many
+from lightly_studio.resolvers.video_frame_resolver.get_all_by_dataset_id import (
+    get_all_by_dataset_id,
+)
+from lightly_studio.resolvers.video_frame_resolver.get_by_id import (
+    get_by_id,
+)
+from lightly_studio.resolvers.video_frame_resolver.get_table_fields_bounds import (
+    get_table_fields_bounds,
+)
+
+__all__ = [
+    "count_video_frames_annotations",
+    "create_many",
+    "get_all_by_dataset_id",
+    "get_by_id",
+    "get_table_fields_bounds",
+]

lightly_studio/resolvers/video_frame_resolver/count_video_frames_annotations.py

@@ -0,0 +1,83 @@
+"""Count video frames annotations."""
+
+from typing import Any, List, Optional, Tuple
+from uuid import UUID
+
+from sqlmodel import Session, asc, col, func, select
+from sqlmodel.sql.expression import Select
+
+from lightly_studio.models.annotation.annotation_base import AnnotationBaseTable
+from lightly_studio.models.annotation_label import AnnotationLabelTable
+from lightly_studio.models.sample import SampleTable
+from lightly_studio.models.video import VideoFrameTable
+from lightly_studio.resolvers.video_frame_resolver.video_frame_annotations_counter_filter import (
+    VideoFrameAnnotationsCounterFilter,
+)
+from lightly_studio.resolvers.video_resolver.count_video_frame_annotations_by_video_dataset import (
+    CountAnnotationsView,
+)
+
+
+def count_video_frames_annotations(
+    session: Session, dataset_id: UUID, filters: Optional[VideoFrameAnnotationsCounterFilter] = None
+) -> List[CountAnnotationsView]:
+    """Count the annotations by video frames."""
+    unfiltered_query = (
+        _build_base_query(dataset_id=dataset_id, count_column_name="total")
+        .group_by(col(AnnotationBaseTable.annotation_label_id))
+        .subquery("unfiltered")
+    )
+
+    filtered_query = _build_base_query(dataset_id=dataset_id, count_column_name="filtered_count")
+
+    if filters:
+        filtered_query = filters.apply(filtered_query)
+
+    filtered_subquery = filtered_query.group_by(
+        col(AnnotationBaseTable.annotation_label_id)
+    ).subquery("filtered")
+
+    final_query: Select[Any] = (
+        select(
+            col(AnnotationLabelTable.annotation_label_name).label("label"),
+            col(unfiltered_query.c.total).label("total"),
+            func.coalesce(filtered_subquery.c.filtered_count, 0).label("filtered_count"),
+        )
+        .select_from(AnnotationLabelTable)
+        .join(
+            unfiltered_query,
+            unfiltered_query.c.label_id == col(AnnotationLabelTable.annotation_label_id),
+        )
+        .outerjoin(
+            filtered_subquery,
+            filtered_subquery.c.label_id == col(AnnotationLabelTable.annotation_label_id),
+        )
+        .order_by(asc(AnnotationLabelTable.annotation_label_name))
+    )
+
+    rows = session.execute(final_query).mappings().all()
+
+    return [
+        CountAnnotationsView(
+            label_name=row["label"],
+            total_count=row["total"],
+            current_count=row["filtered_count"],
+        )
+        for row in rows
+    ]
+
+
+def _build_base_query(dataset_id: UUID, count_column_name: str) -> Select[Tuple[Any, int]]:
+    return (
+        select(
+            col(AnnotationBaseTable.annotation_label_id).label("label_id"),
+            func.count(col(AnnotationBaseTable.annotation_label_id)).label(count_column_name),
+        )
+        .select_from(AnnotationBaseTable)
+        .join(
+            VideoFrameTable,
+            col(VideoFrameTable.sample_id) == col(AnnotationBaseTable.parent_sample_id),
+        )
+        .join(SampleTable, col(VideoFrameTable.parent_sample_id) == col(SampleTable.sample_id))
+        .where(col(SampleTable.dataset_id) == dataset_id)
+    )

lightly_studio/resolvers/video_frame_resolver/create_many.py

@@ -0,0 +1,57 @@
+"""Implementation of create functions for video_frames."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from sqlmodel import Session
+
+from lightly_studio.models.dataset import SampleType
+from lightly_studio.models.sample import SampleCreate
+from lightly_studio.models.video import VideoFrameCreate, VideoFrameTable
+from lightly_studio.resolvers import dataset_resolver, sample_resolver
+
+
+class VideoFrameCreateHelper(VideoFrameCreate):
+    """Helper class to create VideoFrameTable with sample_id."""
+
+    sample_id: UUID
+
+
+def create_many(session: Session, dataset_id: UUID, samples: list[VideoFrameCreate]) -> list[UUID]:
+    """Create multiple video_frame samples in a single database commit.
+
+    Args:
+        session: The database session.
+        dataset_id: The uuid of the dataset to attach to.
+        samples: The video_frames to create in the database.
+
+    Returns:
+        List of UUIDs of VideoFrameTable entries that got added to the database.
+    """
+    dataset_resolver.check_dataset_type(
+        session=session,
+        dataset_id=dataset_id,
+        expected_type=SampleType.VIDEO_FRAME,
+    )
+    sample_ids = sample_resolver.create_many(
+        session=session,
+        samples=[SampleCreate(dataset_id=dataset_id) for _ in samples],
+    )
+    # Bulk create VideoFrameTable entries using the generated sample_ids.
+    db_video_frames = [
+        VideoFrameTable.model_validate(
+            VideoFrameCreateHelper(
+                frame_number=sample.frame_number,
+                frame_timestamp_s=sample.frame_timestamp_s,
+                frame_timestamp_pts=sample.frame_timestamp_pts,
+                parent_sample_id=sample.parent_sample_id,
+                sample_id=sample_id,
+                rotation_deg=sample.rotation_deg,
+            )
+        )
+        for sample_id, sample in zip(sample_ids, samples)
+    ]
+    session.bulk_save_objects(db_video_frames)
+    session.commit()
+    return sample_ids

lightly_studio/resolvers/video_frame_resolver/get_all_by_dataset_id.py

@@ -0,0 +1,63 @@
+"""Implementation of get_all_by_dataset_id function for videos."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any
+from uuid import UUID
+
+from pydantic import BaseModel
+from sqlmodel import Session, col, func, select
+
+from lightly_studio.api.routes.api.validators import Paginated
+from lightly_studio.models.sample import SampleTable
+from lightly_studio.models.video import VideoFrameTable, VideoTable
+from lightly_studio.resolvers.video_frame_resolver.video_frame_filter import VideoFrameFilter
+
+
+class VideoFramesWithCount(BaseModel):
+    """Result of getting all samples."""
+
+    samples: Sequence[VideoFrameTable]
+    total_count: int
+    next_cursor: int | None = None
+
+
+def get_all_by_dataset_id(
+    session: Session,
+    dataset_id: UUID,
+    pagination: Paginated | None = None,
+    video_frame_filter: VideoFrameFilter | None = None,
+) -> VideoFramesWithCount:
+    """Retrieve video frame samples for a specific dataset with optional filtering."""
+    filters: list[Any] = [SampleTable.dataset_id == dataset_id]
+
+    base_query = (
+        select(VideoFrameTable)
+        .join(VideoFrameTable.sample)
+        .join(VideoFrameTable.video)
+        .where(*filters)
+    )
+
+    if video_frame_filter:
+        base_query = video_frame_filter.apply(base_query)
+
+    samples_query = base_query.order_by(
+        col(VideoTable.file_path_abs).asc(), col(VideoFrameTable.frame_number).asc()
+    )
+
+    # Apply pagination if provided
+    if pagination is not None:
+        samples_query = samples_query.offset(pagination.offset).limit(pagination.limit)
+
+    total_count_query = select(func.count()).select_from(base_query.subquery())
+    total_count = session.exec(total_count_query).one()
+    next_cursor = None
+    if pagination and pagination.offset + pagination.limit < total_count:
+        next_cursor = pagination.offset + pagination.limit
+
+    return VideoFramesWithCount(
+        samples=session.exec(samples_query).all(),
+        total_count=total_count,
+        next_cursor=next_cursor,
+    )

lightly_studio/resolvers/video_frame_resolver/get_by_id.py

@@ -0,0 +1,13 @@
+"""Retrieve the video frame by ID resolver implementation."""
+
+from uuid import UUID
+
+from sqlmodel import Session, select
+
+from lightly_studio.models.video import VideoFrameTable
+
+
+def get_by_id(session: Session, sample_id: UUID) -> VideoFrameTable:
+    """Retrieve a single video frame by ID within a dataset."""
+    query = select(VideoFrameTable).where(VideoFrameTable.sample_id == sample_id)
+    return session.exec(query).one()