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

lightly_studio/core/dataset.py

@@ -2,522 +2,394 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
 from pathlib import Path
-from typing import Iterable
+from typing import Iterable, Iterator
 from uuid import UUID
 
-import PIL
 from labelformat.formats import (
     COCOInstanceSegmentationInput,
     COCOObjectDetectionInput,
     YOLOv8ObjectDetectionInput,
 )
-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.api.db import db_manager
-from lightly_studio.models.annotation.annotation_base import AnnotationCreate
-from lightly_studio.models.annotation_label import AnnotationLabelCreate
-from lightly_studio.models.annotation_task import (
-    AnnotationTaskTable,
+from sqlmodel import Session, select
+
+from lightly_studio import db_manager
+from lightly_studio.api import features
+from lightly_studio.core import add_samples
+from lightly_studio.core.dataset_query.dataset_query import DatasetQuery
+from lightly_studio.core.dataset_query.match_expression import MatchExpression
+from lightly_studio.core.dataset_query.order_by import OrderByExpression
+from lightly_studio.core.sample import Sample
+from lightly_studio.dataset import fsspec_lister
+from lightly_studio.dataset.embedding_manager import EmbeddingManagerProvider
+from lightly_studio.metadata import compute_typicality
+from lightly_studio.models.annotation.annotation_base import (
     AnnotationType,
 )
 from lightly_studio.models.dataset import DatasetCreate, DatasetTable
-from lightly_studio.models.sample import SampleCreate, SampleTable
+from lightly_studio.models.sample import SampleTable
 from lightly_studio.resolvers import (
-    annotation_label_resolver,
-    annotation_resolver,
-    annotation_task_resolver,
     dataset_resolver,
+    embedding_model_resolver,
     sample_resolver,
 )
 from lightly_studio.type_definitions import PathLike
 
 # 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
+DEFAULT_DATASET_NAME = "default_dataset"
 
-
-@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
+_SliceType = slice  # to avoid shadowing built-in slice in type annotations
 
 
 class Dataset:
     """A LightlyStudio Dataset.
 
-    Represents a dataset in LightlyStudio.
-
-    Args:
-        name: The name of the dataset. If None, a default name will be assigned.
+    Keeps a reference to the underlying DatasetTable.
     """
 
-    def __init__(self, name: str | None = None) -> None:
+    def __init__(self, dataset: DatasetTable) -> None:
         """Initialize a LightlyStudio Dataset."""
+        self._inner = dataset
+        # TODO(Michal, 09/2025): Do not store the session. Instead, use the
+        # dataset object session.
+        self.session = db_manager.persistent_session()
+
+    @staticmethod
+    def create(name: str | None = None) -> Dataset:
+        """Create a new dataset."""
+        if name is None:
+            name = DEFAULT_DATASET_NAME
+
+        dataset = dataset_resolver.create(
+            session=db_manager.persistent_session(),
+            dataset=DatasetCreate(name=name, directory=""),
+        )
+        return Dataset(dataset=dataset)
+
+    @staticmethod
+    def load(name: str | None = None) -> Dataset:
+        """Load an existing dataset."""
         if name is None:
             name = "default_dataset"
-        self.name = name
-        self.session = db_manager.persistent_session()
-        # Create dataset.
-        self._dataset = dataset_resolver.create(
-            session=self.session,
-            dataset=DatasetCreate(
-                name=self.name,
-                directory="",  # The directory is not used at the moment
-            ),
+
+        dataset = dataset_resolver.get_by_name(session=db_manager.persistent_session(), name=name)
+        if dataset is None:
+            raise ValueError(f"Dataset with name '{name}' not found.")
+
+        return Dataset(dataset=dataset)
+
+    @staticmethod
+    def load_or_create(name: str | None = None) -> Dataset:
+        """Create a new dataset or load an existing one."""
+        if name is None:
+            name = "default_dataset"
+
+        dataset = dataset_resolver.get_by_name(session=db_manager.persistent_session(), name=name)
+        if dataset is None:
+            return Dataset.create(name=name)
+
+        return Dataset(dataset=dataset)
+
+    def __iter__(self) -> Iterator[Sample]:
+        """Iterate over samples in the dataset."""
+        for sample in self.session.exec(
+            select(SampleTable).where(SampleTable.dataset_id == self.dataset_id)
+        ):
+            yield Sample(inner=sample)
+
+    def get_sample(self, sample_id: UUID) -> Sample:
+        """Get a single sample from the dataset by its ID.
+
+        Args:
+            sample_id: The UUID of the sample to retrieve.
+
+        Returns:
+            A single SampleTable object.
+
+        Raises:
+            IndexError: If no sample is found with the given sample_id.
+        """
+        sample = sample_resolver.get_by_id(
+            self.session, dataset_id=self.dataset_id, sample_id=sample_id
         )
 
+        if sample is None:
+            raise IndexError(f"No sample found for sample_id: {sample_id}")
+        return Sample(inner=sample)
+
     @property
     def dataset_id(self) -> UUID:
         """Get the dataset ID."""
-        return self._dataset.dataset_id
+        return self._inner.dataset_id
+
+    @property
+    def name(self) -> str:
+        """Get the dataset name."""
+        return self._inner.name
+
+    def query(self) -> DatasetQuery:
+        """Create a DatasetQuery for this dataset.
+
+        Returns:
+            A DatasetQuery instance for querying samples in this dataset.
+        """
+        return DatasetQuery(dataset=self._inner, session=self.session)
+
+    def match(self, match_expression: MatchExpression) -> DatasetQuery:
+        """Create a query on the dataset and store a field condition for filtering.
+
+        Args:
+            match_expression: Defines the filter.
+
+        Returns:
+            DatasetQuery for method chaining.
+        """
+        return self.query().match(match_expression)
+
+    def order_by(self, *order_by: OrderByExpression) -> DatasetQuery:
+        """Create a query on the dataset and store ordering expressions.
+
+        Args:
+            order_by: One or more ordering expressions. They are applied in order.
+                E.g. first ordering by sample width and then by sample file_name will
+                only order the samples with the same sample width by file_name.
+
+        Returns:
+            DatasetQuery for method chaining.
+        """
+        return self.query().order_by(*order_by)
+
+    def slice(self, offset: int = 0, limit: int | None = None) -> DatasetQuery:
+        """Create a query on the dataset and apply offset and limit to results.
+
+        Args:
+            offset: Number of items to skip from beginning (default: 0).
+            limit: Maximum number of items to return (None = no limit).
+
+        Returns:
+            DatasetQuery for method chaining.
+        """
+        return self.query().slice(offset, limit)
+
+    def __getitem__(self, key: _SliceType) -> DatasetQuery:
+        """Create a query on the dataset and enable bracket notation for slicing.
+
+        Args:
+            key: A slice object (e.g., [10:20], [:50], [100:]).
+
+        Returns:
+            DatasetQuery with slice applied.
+
+        Raises:
+            TypeError: If key is not a slice object.
+            ValueError: If slice contains unsupported features or conflicts with existing slice.
+        """
+        return self.query()[key]
 
     def add_samples_from_path(
         self,
         path: PathLike,
-        recursive: bool = True,
-        allowed_extensions: Iterable[str] = {
-            ".png",
-            ".jpg",
-            ".jpeg",
-            ".gif",
-            ".webp",
-            ".bmp",
-            ".tiff",
-        },
+        allowed_extensions: Iterable[str] | None = None,
+        embed: bool = True,
     ) -> None:
         """Adding samples from the specified path to the dataset.
 
         Args:
             path: Path to the folder containing the images to add.
-            recursive: If True, search for images recursively in subfolders.
             allowed_extensions: An iterable container of allowed image file
                 extensions.
+            embed: If True, generate embeddings for the newly added samples.
         """
-        path = Path(path).absolute() if isinstance(path, str) else path.absolute()
-        if not path.exists() or not path.is_dir():
-            raise ValueError(f"Provided path is not a valid directory: {path}")
-
         # Collect image file paths.
-        allowed_extensions_set = {ext.lower() for ext in allowed_extensions}
-        image_paths = []
-        path_iter = path.rglob("*") if recursive else path.glob("*")
-        for path in path_iter:
-            if path.is_file() and path.suffix.lower() in allowed_extensions_set:
-                image_paths.append(path)
+        if allowed_extensions:
+            allowed_extensions_set = {ext.lower() for ext in allowed_extensions}
+        else:
+            allowed_extensions_set = None
+        image_paths = list(
+            fsspec_lister.iter_files_from_path(
+                path=str(path), allowed_extensions=allowed_extensions_set
+            )
+        )
         print(f"Found {len(image_paths)} images in {path}.")
 
         # Process images.
-        _load_into_dataset_from_paths(
+        created_sample_ids = add_samples.load_into_dataset_from_paths(
             session=self.session,
             dataset_id=self.dataset_id,
             image_paths=image_paths,
         )
 
+        if embed:
+            _generate_embeddings(
+                session=self.session, dataset_id=self.dataset_id, sample_ids=created_sample_ids
+            )
+
     def add_samples_from_labelformat(
         self,
         input_labels: ObjectDetectionInput | InstanceSegmentationInput,
         images_path: PathLike,
-        is_prediction: bool = True,
-        task_name: str | None = None,
+        embed: bool = True,
     ) -> None:
         """Load a dataset from a labelformat object and store in database.
 
         Args:
             input_labels: The labelformat input object.
             images_path: Path to the folder containing the images.
-            is_prediction: Whether the task is for prediction or labels.
-            task_name: Optional name for the annotation task. If None, a
-                default name is generated.
-
-        Returns:
-            DatasetTable: The created dataset table entry.
+            embed: If True, generate embeddings for the newly added samples.
         """
         if isinstance(images_path, str):
             images_path = Path(images_path)
         images_path = images_path.absolute()
 
-        # Determine annotation type based on input.
-        # Currently, we always create BBOX tasks, even for segmentation,
-        # as segmentation data is stored alongside bounding boxes.
-        annotation_type = AnnotationType.BBOX
-
-        # Generate a default task name if none is provided.
-        if task_name is None:
-            task_name = f"Loaded from labelformat: {self.name}"
-
-        # Create annotation task.
-        new_annotation_task = annotation_task_resolver.create(
-            session=self.session,
-            annotation_task=AnnotationTaskTable(
-                name=task_name,
-                annotation_type=annotation_type,
-                is_prediction=is_prediction,
-            ),
-        )
-
-        _load_into_dataset(
+        created_sample_ids = add_samples.load_into_dataset_from_labelformat(
             session=self.session,
             dataset_id=self.dataset_id,
             input_labels=input_labels,
             images_path=images_path,
-            annotation_task_id=new_annotation_task.annotation_task_id,
         )
 
-    def from_yolo(
+        if embed:
+            _generate_embeddings(
+                session=self.session, dataset_id=self.dataset_id, sample_ids=created_sample_ids
+            )
+
+    def add_samples_from_yolo(
         self,
-        data_yaml_path: str,
+        data_yaml: PathLike,
         input_split: str = "train",
-        task_name: str | None = None,
-    ) -> DatasetTable:
+        embed: bool = True,
+    ) -> None:
         """Load a dataset in YOLO format and store in DB.
 
         Args:
-            data_yaml_path: Path to the YOLO data.yaml file.
+            data_yaml: Path to the YOLO data.yaml file.
             input_split: The split to load (e.g., 'train', 'val').
-            task_name: Optional name for the annotation task. If None, a
-                default name is generated.
-
-        Returns:
-            DatasetTable: The created dataset table entry.
+            embed: If True, generate embeddings for the newly added samples.
         """
-        data_yaml = Path(data_yaml_path).absolute()
-        dataset_name = data_yaml.parent.name
+        if isinstance(data_yaml, str):
+            data_yaml = Path(data_yaml)
+        data_yaml = data_yaml.absolute()
 
-        if task_name is None:
-            task_name = f"Loaded from YOLO: {data_yaml.name} ({input_split} split)"
+        if not data_yaml.is_file() or data_yaml.suffix != ".yaml":
+            raise FileNotFoundError(f"YOLO data yaml file not found: '{data_yaml}'")
 
         # Load the dataset using labelformat.
         label_input = YOLOv8ObjectDetectionInput(
             input_file=data_yaml,
             input_split=input_split,
         )
-        img_dir = label_input._images_dir()  # noqa: SLF001
+        images_path = label_input._images_dir()  # noqa: SLF001
 
-        return self.from_labelformat(  # type: ignore[no-any-return,attr-defined]
+        self.add_samples_from_labelformat(
             input_labels=label_input,
-            dataset_name=dataset_name,
-            img_dir=str(img_dir),
-            is_prediction=False,
-            task_name=task_name,
+            images_path=images_path,
+            embed=embed,
         )
 
-    def from_coco_object_detections(
+    def add_samples_from_coco(
         self,
-        annotations_json_path: str,
-        img_dir: str,
-        task_name: str | None = None,
-    ) -> DatasetTable:
+        annotations_json: PathLike,
+        images_path: PathLike,
+        annotation_type: AnnotationType = AnnotationType.OBJECT_DETECTION,
+        embed: bool = True,
+    ) -> None:
         """Load a dataset in COCO Object Detection format and store in DB.
 
         Args:
-            annotations_json_path: Path to the COCO annotations JSON file.
-            img_dir: Path to the folder containing the images.
-            task_name: Optional name for the annotation task. If None, a
-                default name is generated.
-
-        Returns:
-            DatasetTable: The created dataset table entry.
+            annotations_json: Path to the COCO annotations JSON file.
+            images_path: Path to the folder containing the images.
+            annotation_type: The type of annotation to be loaded (e.g., 'ObjectDetection',
+                'InstanceSegmentation').
+            embed: If True, generate embeddings for the newly added samples.
         """
-        annotations_json = Path(annotations_json_path)
-        dataset_name = annotations_json.parent.name
-
-        if task_name is None:
-            task_name = f"Loaded from COCO Object Detection: {annotations_json.name}"
+        if isinstance(annotations_json, str):
+            annotations_json = Path(annotations_json)
+        annotations_json = annotations_json.absolute()
 
-        label_input = COCOObjectDetectionInput(
-            input_file=annotations_json,
-        )
-        img_dir_path = Path(img_dir).absolute()
-
-        return self.from_labelformat(  # type: ignore[no-any-return, attr-defined]
-            input_labels=label_input,
-            dataset_name=dataset_name,
-            img_dir=str(img_dir_path),
-            is_prediction=False,
-            task_name=task_name,
-        )
+        if not annotations_json.is_file() or annotations_json.suffix != ".json":
+            raise FileNotFoundError(f"COCO annotations json file not found: '{annotations_json}'")
 
-    def from_coco_instance_segmentations(
-        self,
-        annotations_json_path: str,
-        img_dir: str,
-        task_name: str | None = None,
-    ) -> DatasetTable:
-        """Load a dataset in COCO Instance Segmentation format and store in DB.
-
-        Args:
-            annotations_json_path: Path to the COCO annotations JSON file.
-            img_dir: Path to the folder containing the images.
-            task_name: Optional name for the annotation task. If None, a
-                default name is generated.
+        label_input: COCOObjectDetectionInput | COCOInstanceSegmentationInput
 
-        Returns:
-            DatasetTable: The created dataset table entry.
-        """
-        annotations_json = Path(annotations_json_path)
-        dataset_name = annotations_json.parent.name
-
-        if task_name is None:
-            task_name = f"Loaded from COCO Instance Segmentation: {annotations_json.name}"
+        if annotation_type == AnnotationType.OBJECT_DETECTION:
+            label_input = COCOObjectDetectionInput(
+                input_file=annotations_json,
+            )
+        elif annotation_type == AnnotationType.INSTANCE_SEGMENTATION:
+            label_input = COCOInstanceSegmentationInput(
+                input_file=annotations_json,
+            )
+        else:
+            raise ValueError(f"Invalid annotation type: {annotation_type}")
 
-        label_input = COCOInstanceSegmentationInput(
-            input_file=annotations_json,
-        )
-        img_dir_path = Path(img_dir).absolute()
+        images_path = Path(images_path).absolute()
 
-        return self.from_labelformat(  # type: ignore[no-any-return,attr-defined]
+        self.add_samples_from_labelformat(
             input_labels=label_input,
-            dataset_name=dataset_name,
-            img_dir=str(img_dir_path),
-            is_prediction=False,
-            task_name=task_name,
+            images_path=images_path,
+            embed=embed,
         )
 
-    @staticmethod
-    def load_from_db(name: str, db_path: PathLike) -> Dataset:
-        """Load a dataset from the database.
+    def compute_typicality_metadata(
+        self,
+        embedding_model_name: str | None = None,
+        metadata_name: str = "typicality",
+    ) -> None:
+        """Computes typicality from embeddings, for K nearest neighbors.
 
-        Returns:
-            Dataset: The loaded dataset.
+        Args:
+            embedding_model_name:
+                The name of the embedding model to use. If not given, the default
+                embedding model is used.
+            metadata_name:
+                The name of the metadata to store the typicality values in. If not give, the default
+                name "typicality" is used.
         """
-        raise NotImplementedError
-
-
-def _load_into_dataset_from_paths(
-    dataset_id: UUID,
-    session: Session,
-    image_paths: Iterable[Path],
-) -> None:
-    samples_to_create: list[SampleCreate] = []
-
-    for image_path in tqdm(
-        image_paths,
-        desc="Processing images",
-        unit=" images",
-    ):
-        try:
-            image = PIL.Image.open(image_path)
-            width, height = image.size
-            image.close()
-        except (FileNotFoundError, PIL.UnidentifiedImageError, OSError):
-            continue
-
-        sample = SampleCreate(
-            file_name=image_path.name,
-            file_path_abs=str(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:
-            _ = sample_resolver.create_many(session=session, samples=samples_to_create)
-            samples_to_create = []
-
-    # Handle remaining samples
-    if samples_to_create:
-        _ = sample_resolver.create_many(session=session, samples=samples_to_create)
-
-
-def _load_into_dataset(
-    session: Session,
-    dataset_id: UUID,
-    input_labels: ObjectDetectionInput | InstanceSegmentationInput,
-    images_path: Path,
-    annotation_task_id: UUID,
-) -> None:
-    """Store a loaded dataset in database."""
-    # Create label mapping
-    label_map = _create_label_map(session=session, input_labels=input_labels)
-
-    annotations_to_create: list[AnnotationCreate] = []
-    sample_ids: list[UUID] = []
-    samples_to_create: list[SampleCreate] = []
-    samples_image_data: list[
-        tuple[SampleCreate, 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)
-        samples_image_data.append((sample, typed_image_data))
-
-        if len(samples_to_create) >= SAMPLE_BATCH_SIZE:
-            stored_samples = sample_resolver.create_many(session=session, samples=samples_to_create)
-            _process_batch_annotations(
-                session=session,
-                stored_samples=stored_samples,
-                samples_data=samples_image_data,
-                dataset_id=dataset_id,
-                label_map=label_map,
-                annotation_task_id=annotation_task_id,
-                annotations_to_create=annotations_to_create,
-                sample_ids=sample_ids,
-            )
-            samples_to_create.clear()
-            samples_image_data.clear()
-
-    if samples_to_create:
-        stored_samples = sample_resolver.create_many(session=session, samples=samples_to_create)
-        _process_batch_annotations(
-            session=session,
-            stored_samples=stored_samples,
-            samples_data=samples_image_data,
-            dataset_id=dataset_id,
-            label_map=label_map,
-            annotation_task_id=annotation_task_id,
-            annotations_to_create=annotations_to_create,
-            sample_ids=sample_ids,
-        )
-
-    # Insert any remaining annotations
-    if annotations_to_create:
-        annotation_resolver.create_many(session=session, annotations=annotations_to_create)
-
-
-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,
-    image_data: ImageObjectDetection,
-) -> list[AnnotationCreate]:
-    """Process object detection annotations for a single image."""
-    new_annotations = []
-    for obj in image_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=x,
-                y=y,
-                width=width,
-                height=height,
-                confidence=obj.confidence,
-                annotation_task_id=context.annotation_task_id,
-            )
-        )
-    return new_annotations
-
-
-def _process_instance_segmentation_annotations(
-    context: AnnotationProcessingContext,
-    image_data: ImageInstanceSegmentation,
-) -> list[AnnotationCreate]:
-    """Process instance segmentation annotations for a single image."""
-    new_annotations = []
-    for obj in image_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=x,
-                y=y,
-                width=width,
-                height=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],
-    samples_data: list[tuple[SampleCreate, ImageInstanceSegmentation | ImageObjectDetection]],
-    dataset_id: UUID,
-    label_map: dict[int, UUID],
-    annotation_task_id: UUID,
-    annotations_to_create: list[AnnotationCreate],
-    sample_ids: list[UUID],
-) -> None:
-    """Process annotations for a batch of samples."""
-    for stored_sample, (_, img_data) in zip(stored_samples, samples_data):
-        sample_ids.append(stored_sample.sample_id)
-
-        context = AnnotationProcessingContext(
-            dataset_id=dataset_id,
-            sample_id=stored_sample.sample_id,
-            label_map=label_map,
-            annotation_task_id=annotation_task_id,
+        embedding_model_id = embedding_model_resolver.get_by_name(
+            session=self.session,
+            dataset_id=self.dataset_id,
+            embedding_model_name=embedding_model_name,
+        ).embedding_model_id
+        compute_typicality.compute_typicality_metadata(
+            session=self.session,
+            dataset_id=self.dataset_id,
+            embedding_model_id=embedding_model_id,
+            metadata_name=metadata_name,
         )
 
-        if isinstance(img_data, ImageInstanceSegmentation):
-            new_annotations = _process_instance_segmentation_annotations(
-                context=context, image_data=img_data
-            )
-        elif isinstance(img_data, ImageObjectDetection):
-            new_annotations = _process_object_detection_annotations(
-                context=context, image_data=img_data
-            )
-        else:
-            raise ValueError(f"Unsupported annotation type: {type(img_data)}")
 
-        annotations_to_create.extend(new_annotations)
+def _generate_embeddings(session: Session, dataset_id: UUID, sample_ids: list[UUID]) -> None:
+    """Generate and store embeddings for samples.
 
-        if len(annotations_to_create) >= ANNOTATION_BATCH_SIZE:
-            annotation_resolver.create_many(session=session, annotations=annotations_to_create)
-            annotations_to_create.clear()
+    Args:
+        session: Database session for resolver operations.
+        dataset_id: The ID of the dataset to associate with the embedding model.
+        sample_ids: List of sample IDs to generate embeddings for.
+    """
+    if not sample_ids:
+        return
+
+    embedding_manager = EmbeddingManagerProvider.get_embedding_manager()
+    model_id = embedding_manager.load_or_get_default_model(
+        session=session,
+        dataset_id=dataset_id,
+    )
+    if model_id is None:
+        print("No embedding model loaded. Skipping embedding generation.")
+        return
+
+    embedding_manager.embed_images(
+        session=session,
+        sample_ids=sample_ids,
+        embedding_model_id=model_id,
+    )
+
+    # Mark the embedding search feature as enabled.
+    if "embeddingSearchEnabled" not in features.lightly_studio_active_features:
+        features.lightly_studio_active_features.append("embeddingSearchEnabled")