lightly-studio 0.3.1__py3-none-any.whl → 0.3.2__py3-none-any.whl

lightly_studio/core/add_samples.py

@@ -0,0 +1,383 @@
+"""Functions to add samples and their annotations to a dataset in the database."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Iterable
+from uuid import UUID
+
+import fsspec
+import PIL
+from labelformat.model.binary_mask_segmentation import BinaryMaskSegmentation
+from labelformat.model.bounding_box import BoundingBoxFormat
+from labelformat.model.image import Image
+from labelformat.model.instance_segmentation import (
+    ImageInstanceSegmentation,
+    InstanceSegmentationInput,
+)
+from labelformat.model.multipolygon import MultiPolygon
+from labelformat.model.object_detection import (
+    ImageObjectDetection,
+    ObjectDetectionInput,
+)
+from sqlmodel import Session
+from tqdm import tqdm
+
+from lightly_studio.models.annotation.annotation_base import AnnotationCreate
+from lightly_studio.models.annotation_label import AnnotationLabelCreate
+from lightly_studio.models.sample import SampleCreate, SampleTable
+from lightly_studio.resolvers import (
+    annotation_label_resolver,
+    annotation_resolver,
+    sample_resolver,
+)
+
+# Constants
+ANNOTATION_BATCH_SIZE = 64  # Number of annotations to process in a single batch
+SAMPLE_BATCH_SIZE = 32  # Number of samples to process in a single batch
+MAX_EXAMPLE_PATHS_TO_SHOW = 5
+
+
+@dataclass
+class _AnnotationProcessingContext:
+    """Context for processing annotations for a single sample."""
+
+    dataset_id: UUID
+    sample_id: UUID
+    label_map: dict[int, UUID]
+    annotation_task_id: UUID
+
+
+@dataclass
+class _LoadingLoggingContext:
+    """Context for the logging while loading data."""
+
+    n_samples_before_loading: int
+    n_samples_to_be_inserted: int = 0
+    example_paths_not_inserted: list[str] = field(default_factory=list)
+
+    def update_example_paths(self, example_paths_not_inserted: list[str]) -> None:
+        if len(self.example_paths_not_inserted) >= MAX_EXAMPLE_PATHS_TO_SHOW:
+            return
+        upper_limit = MAX_EXAMPLE_PATHS_TO_SHOW - len(self.example_paths_not_inserted)
+        self.example_paths_not_inserted.extend(example_paths_not_inserted[:upper_limit])
+
+
+def load_into_dataset_from_paths(
+    session: Session,
+    dataset_id: UUID,
+    image_paths: Iterable[str],
+) -> list[UUID]:
+    """Load images from file paths into the dataset.
+
+    Args:
+        session: The database session.
+        dataset_id: The ID of the dataset to load images into.
+        image_paths: An iterable of file paths to the images to load.
+
+    Returns:
+        A list of UUIDs of the created samples.
+    """
+    samples_to_create: list[SampleCreate] = []
+    created_sample_ids: list[UUID] = []
+
+    logging_context = _LoadingLoggingContext(
+        n_samples_to_be_inserted=sum(1 for _ in image_paths),
+        n_samples_before_loading=sample_resolver.count_by_dataset_id(
+            session=session, dataset_id=dataset_id
+        ),
+    )
+
+    for image_path in tqdm(
+        image_paths,
+        desc="Processing images",
+        unit=" images",
+    ):
+        try:
+            with fsspec.open(image_path, "rb") as file:
+                image = PIL.Image.open(file)
+                width, height = image.size
+                image.close()
+        except (FileNotFoundError, PIL.UnidentifiedImageError, OSError):
+            continue
+
+        sample = SampleCreate(
+            file_name=Path(image_path).name,
+            file_path_abs=image_path,
+            width=width,
+            height=height,
+            dataset_id=dataset_id,
+        )
+        samples_to_create.append(sample)
+
+        # Process batch when it reaches SAMPLE_BATCH_SIZE
+        if len(samples_to_create) >= SAMPLE_BATCH_SIZE:
+            created_samples_batch, paths_not_inserted = _create_batch_samples(
+                session=session, samples=samples_to_create
+            )
+            created_sample_ids.extend(s.sample_id for s in created_samples_batch)
+            logging_context.update_example_paths(paths_not_inserted)
+            samples_to_create = []
+
+    # Handle remaining samples
+    if samples_to_create:
+        created_samples_batch, paths_not_inserted = _create_batch_samples(
+            session=session, samples=samples_to_create
+        )
+        created_sample_ids.extend(s.sample_id for s in created_samples_batch)
+        logging_context.update_example_paths(paths_not_inserted)
+
+    _log_loading_results(session=session, dataset_id=dataset_id, logging_context=logging_context)
+    return created_sample_ids
+
+
+def load_into_dataset_from_labelformat(
+    session: Session,
+    dataset_id: UUID,
+    input_labels: ObjectDetectionInput | InstanceSegmentationInput,
+    images_path: Path,
+    annotation_task_id: UUID,
+) -> list[UUID]:
+    """Load samples and their annotations from a labelformat input into the dataset.
+
+    Args:
+        session: The database session.
+        dataset_id: The ID of the dataset to load samples into.
+        input_labels: The labelformat input containing images and annotations.
+        images_path: The path to the directory containing the images.
+        annotation_task_id: The ID of the annotation task to associate with the annotations.
+
+    Returns:
+        A list of UUIDs of the created samples.
+    """
+    logging_context = _LoadingLoggingContext(
+        n_samples_to_be_inserted=sum(1 for _ in input_labels.get_labels()),
+        n_samples_before_loading=sample_resolver.count_by_dataset_id(
+            session=session, dataset_id=dataset_id
+        ),
+    )
+
+    # Create label mapping
+    label_map = _create_label_map(session=session, input_labels=input_labels)
+
+    annotations_to_create: list[AnnotationCreate] = []
+    samples_to_create: list[SampleCreate] = []
+    created_sample_ids: list[UUID] = []
+    image_path_to_anno_data: dict[str, ImageInstanceSegmentation | ImageObjectDetection] = {}
+
+    for image_data in tqdm(input_labels.get_labels(), desc="Processing images", unit=" images"):
+        image: Image = image_data.image  # type: ignore[attr-defined]
+
+        typed_image_data: ImageInstanceSegmentation | ImageObjectDetection = image_data  # type: ignore[assignment]
+        sample = SampleCreate(
+            file_name=str(image.filename),
+            file_path_abs=str(images_path / image.filename),
+            width=image.width,
+            height=image.height,
+            dataset_id=dataset_id,
+        )
+        samples_to_create.append(sample)
+        image_path_to_anno_data[sample.file_path_abs] = typed_image_data
+
+        if len(samples_to_create) >= SAMPLE_BATCH_SIZE:
+            created_samples_batch, paths_not_inserted = _create_batch_samples(
+                session=session, samples=samples_to_create
+            )
+            created_sample_ids.extend(s.sample_id for s in created_samples_batch)
+            logging_context.update_example_paths(paths_not_inserted)
+            _process_batch_annotations(
+                session=session,
+                stored_samples=created_samples_batch,
+                image_path_to_anno_data=image_path_to_anno_data,
+                dataset_id=dataset_id,
+                label_map=label_map,
+                annotation_task_id=annotation_task_id,
+                annotations_to_create=annotations_to_create,
+            )
+            samples_to_create.clear()
+            image_path_to_anno_data.clear()
+
+    if samples_to_create:
+        created_samples_batch, paths_not_inserted = _create_batch_samples(
+            session=session, samples=samples_to_create
+        )
+        created_sample_ids.extend(s.sample_id for s in created_samples_batch)
+        logging_context.update_example_paths(paths_not_inserted)
+        _process_batch_annotations(
+            session=session,
+            stored_samples=created_samples_batch,
+            image_path_to_anno_data=image_path_to_anno_data,
+            dataset_id=dataset_id,
+            label_map=label_map,
+            annotation_task_id=annotation_task_id,
+            annotations_to_create=annotations_to_create,
+        )
+
+    # Insert any remaining annotations
+    if annotations_to_create:
+        annotation_resolver.create_many(session=session, annotations=annotations_to_create)
+
+    _log_loading_results(session=session, dataset_id=dataset_id, logging_context=logging_context)
+
+    return created_sample_ids
+
+
+def _log_loading_results(
+    session: Session, dataset_id: UUID, logging_context: _LoadingLoggingContext
+) -> None:
+    n_samples_end = sample_resolver.count_by_dataset_id(session=session, dataset_id=dataset_id)
+    n_samples_inserted = n_samples_end - logging_context.n_samples_before_loading
+    print(
+        f"Added {n_samples_inserted} out of {logging_context.n_samples_to_be_inserted}"
+        " new samples to the dataset."
+    )
+    if logging_context.example_paths_not_inserted:
+        # TODO(Jonas, 09/2025): Use logging instead of print
+        print(
+            f"Examples of paths that were not added: "
+            f" {', '.join(logging_context.example_paths_not_inserted)}"
+        )
+
+
+def _create_batch_samples(
+    session: Session, samples: list[SampleCreate]
+) -> tuple[list[SampleTable], list[str]]:
+    """Create the batch samples.
+
+    Args:
+        session: The database session.
+        samples: The samples to create.
+
+    Returns:
+        created_samples: A list of created SampleTable objects,
+        existing_file_paths: A list of file paths that already existed in the database,
+    """
+    file_paths_abs_mapping = {sample.file_path_abs: sample for sample in samples}
+    file_paths_new, file_paths_exist = sample_resolver.filter_new_paths(
+        session=session, file_paths_abs=list(file_paths_abs_mapping.keys())
+    )
+    samples_to_create_filtered = [
+        file_paths_abs_mapping[file_path_new] for file_path_new in file_paths_new
+    ]
+    return (
+        sample_resolver.create_many(session=session, samples=samples_to_create_filtered),
+        file_paths_exist,
+    )
+
+
+def _create_label_map(
+    session: Session,
+    input_labels: ObjectDetectionInput | InstanceSegmentationInput,
+) -> dict[int, UUID]:
+    """Create a mapping of category IDs to annotation label IDs."""
+    label_map = {}
+    for category in tqdm(
+        input_labels.get_categories(),
+        desc="Processing categories",
+        unit=" categories",
+    ):
+        label = AnnotationLabelCreate(annotation_label_name=category.name)
+        stored_label = annotation_label_resolver.create(session=session, label=label)
+        label_map[category.id] = stored_label.annotation_label_id
+    return label_map
+
+
+def _process_object_detection_annotations(
+    context: _AnnotationProcessingContext,
+    anno_data: ImageObjectDetection,
+) -> list[AnnotationCreate]:
+    """Process object detection annotations for a single image."""
+    new_annotations = []
+    for obj in anno_data.objects:
+        box = obj.box.to_format(BoundingBoxFormat.XYWH)
+        x, y, width, height = box
+
+        new_annotations.append(
+            AnnotationCreate(
+                dataset_id=context.dataset_id,
+                sample_id=context.sample_id,
+                annotation_label_id=context.label_map[obj.category.id],
+                annotation_type="object_detection",
+                x=int(x),
+                y=int(y),
+                width=int(width),
+                height=int(height),
+                confidence=obj.confidence,
+                annotation_task_id=context.annotation_task_id,
+            )
+        )
+    return new_annotations
+
+
+def _process_instance_segmentation_annotations(
+    context: _AnnotationProcessingContext,
+    anno_data: ImageInstanceSegmentation,
+) -> list[AnnotationCreate]:
+    """Process instance segmentation annotations for a single image."""
+    new_annotations = []
+    for obj in anno_data.objects:
+        segmentation_rle: None | list[int] = None
+        if isinstance(obj.segmentation, MultiPolygon):
+            box = obj.segmentation.bounding_box().to_format(BoundingBoxFormat.XYWH)
+        elif isinstance(obj.segmentation, BinaryMaskSegmentation):
+            box = obj.segmentation.bounding_box.to_format(BoundingBoxFormat.XYWH)
+            segmentation_rle = obj.segmentation._rle_row_wise  # noqa: SLF001
+        else:
+            raise ValueError(f"Unsupported segmentation type: {type(obj.segmentation)}")
+
+        x, y, width, height = box
+
+        new_annotations.append(
+            AnnotationCreate(
+                dataset_id=context.dataset_id,
+                sample_id=context.sample_id,
+                annotation_label_id=context.label_map[obj.category.id],
+                annotation_type="instance_segmentation",
+                x=int(x),
+                y=int(y),
+                width=int(width),
+                height=int(height),
+                segmentation_mask=segmentation_rle,
+                annotation_task_id=context.annotation_task_id,
+            )
+        )
+    return new_annotations
+
+
+def _process_batch_annotations(  # noqa: PLR0913
+    session: Session,
+    stored_samples: list[SampleTable],
+    image_path_to_anno_data: dict[str, ImageInstanceSegmentation | ImageObjectDetection],
+    dataset_id: UUID,
+    label_map: dict[int, UUID],
+    annotation_task_id: UUID,
+    annotations_to_create: list[AnnotationCreate],
+) -> None:
+    """Process annotations for a batch of samples."""
+    for stored_sample in stored_samples:
+        anno_data = image_path_to_anno_data[stored_sample.file_path_abs]
+
+        context = _AnnotationProcessingContext(
+            dataset_id=dataset_id,
+            sample_id=stored_sample.sample_id,
+            label_map=label_map,
+            annotation_task_id=annotation_task_id,
+        )
+
+        if isinstance(anno_data, ImageInstanceSegmentation):
+            new_annotations = _process_instance_segmentation_annotations(
+                context=context, anno_data=anno_data
+            )
+        elif isinstance(anno_data, ImageObjectDetection):
+            new_annotations = _process_object_detection_annotations(
+                context=context, anno_data=anno_data
+            )
+        else:
+            raise ValueError(f"Unsupported annotation type: {type(anno_data)}")
+
+        annotations_to_create.extend(new_annotations)
+
+        if len(annotations_to_create) >= ANNOTATION_BATCH_SIZE:
+            annotation_resolver.create_many(session=session, annotations=annotations_to_create)
+            annotations_to_create.clear()