Finishing implementation of slicing a VolumeDataset

Author: Lucashsmello

datamint/dataset/annotation_processor.py

@@ -53,10 +53,12 @@ def __init__(
         seglabel2code: dict[str, int],
         image_labels_set: list[str],
         image_lcodes: dict[str, dict[str, int]],
+        allow_external_annotations: bool = False,
     ):
         self.seglabel2code = seglabel2code
         self.image_labels_set = image_labels_set
         self.image_lcodes = image_lcodes
+        self.allow_external_annotations = allow_external_annotations
 
     def collate_frame_segmentations(self,
                                     fr_anns: Sequence['Annotation'],
@@ -67,7 +69,15 @@ def collate_frame_segmentations(self,
         for ann in fr_anns:
             try:
                 seg = self.load_segmentation_data(ann)
-                seg_code_i = self.seglabel2code.get(ann.identifier, 0)
+                seg_code_i = self.seglabel2code.get(ann.identifier)
+                if seg_code_i is None:
+                    if self.allow_external_annotations and ann.identifier:
+                        seg_code_i = max(self.seglabel2code.values(), default=0) + 1
+                        self.seglabel2code[ann.identifier] = seg_code_i
+                        _LOGGER.info(f"Dynamically added segmentation label '{ann.identifier}' with code {seg_code_i}")
+                    else:
+                        raise ValueError(f"Unknown segmentation label '{ann.identifier}' and external annotations are not allowed")
+
                 if seg_code != -1 and seg_code != seg_code_i:
                     raise ValueError(f"Conflicting segmentation codes for frame annotations: "
                                      f"{seg_code} vs {seg_code_i}")
@@ -145,8 +155,17 @@ def load_image_segmentations(self,
             author = ann.created_by or ann.created_by_model or "unknown"
 
             try:
+                seg_code = self.seglabel2code.get(ann.identifier)
+                if seg_code is None:
+                    if self.allow_external_annotations:
+                        seg_code = max(self.seglabel2code.values(), default=0) + 1
+                        self.seglabel2code[ann.identifier] = seg_code
+                        _LOGGER.info(f"Dynamically added segmentation label '{ann.identifier}' with code {seg_code}")
+                    else:
+                        raise ValueError(f"Segmentation annotation {ann.id} has unknown identifier "
+                                         f"{ann.identifier} with no corresponding code in {self.seglabel2code=}")
                 seg = self.load_segmentation_data(ann)
-                seg_code = self.seglabel2code.get(ann.identifier, 0)
+
                 # seg shape: (#slices, H, W)
             except Exception as e:
                 _LOGGER.error(f"Failed to load segmentation for annotation {ann.id}: {e}")
@@ -331,9 +350,20 @@ def convert_image_labels(
         Returns:
             Dict of annotator_id -> one-hot tensor of shape (num_labels,).
         """
-        labels_ret_size = (len(self.image_labels_set),)
         label2code = self.image_lcodes.get('multilabel', {})
 
+        # If allow_external_annotations, first pass to discover unknown labels
+        if self.allow_external_annotations:
+            for ann in annotations:
+                if ann.annotation_type != 'label':
+                    continue
+                if ann.identifier not in label2code:
+                    new_code = len(self.image_labels_set)
+                    self.image_labels_set.append(ann.identifier)
+                    label2code[ann.identifier] = new_code
+                    _LOGGER.info(f"Dynamically added image label '{ann.identifier}' with code {new_code}")
+
+        labels_ret_size = (len(self.image_labels_set),)
         labels_by_user: dict[str, torch.Tensor] = {}
 
         for ann in annotations:
@@ -355,44 +385,36 @@ def apply_merge_strategy(
         self,
         segmentations: dict[str, Tensor],
         strategy: MergeStrategy,
-        output_shape: tuple[int, ...] | None = None,
     ) -> Tensor: ...
 
     @overload
     def apply_merge_strategy(
         self,
         segmentations: dict[str, np.ndarray],
         strategy: MergeStrategy,
-        output_shape: tuple[int, ...] | None = None,
     ) -> np.ndarray: ...
 
     def apply_merge_strategy(
         self,
         segmentations: dict[str, Tensor] | dict[str, np.ndarray],
         strategy: MergeStrategy,
-        output_shape: tuple[int, ...] | None = None,
     ) -> Tensor | np.ndarray:
         """Merge semantic segmentations from multiple annotators.
 
         Args:
             segmentations: Dict of author -> semantic segmentation tensor.
-            output_shape: Shape for empty result if no segmentations are present.
             strategy: Merge strategy ('union', 'intersection', 'mode').
 
         Returns:
             Merged tensor if strategy is specified, otherwise original dict.
         """
         if len(segmentations) == 0:
-            if output_shape is None:
-                raise ValueError("output_shape must be provided when no segmentations are present")
-            empty_segs = torch.zeros(output_shape, dtype=torch.get_default_dtype())
-            empty_segs[0] = 1  # background
-            return empty_segs
+            raise ValueError("No segmentations to merge")
 
         if isinstance(next(iter(segmentations.values())), np.ndarray):
             with torch.no_grad():
                 segmentations = {author: torch.from_numpy(seg) for author, seg in segmentations.items()}
-                return self.apply_merge_strategy(segmentations, strategy, output_shape).numpy()
+                return self.apply_merge_strategy(segmentations, strategy).numpy()
 
         if strategy == 'union':
             merged = self._merge_union(segmentations)

datamint/dataset/base.py

@@ -15,12 +15,11 @@
 import numpy as np
 from datamint.apihandler.dto.annotation_dto import AnnotationType
 from datamint.exceptions import DatamintException
-from datamint.entities import Annotation
 from .annotation_processor import AnnotationProcessor, MergeStrategy
 from datamint.entities.annotations.annotation_spec import AnnotationSpec, CategoryAnnotationSpec
 
 if TYPE_CHECKING:
-    from datamint.entities import Resource, Project
+    from datamint.entities import Resource, Project, Annotation
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -61,10 +60,14 @@ class DatamintBaseDataset(ABC):
         exclude_image_label_names: Blacklist of image labels.
         include_frame_label_names: Whitelist of frame labels.
         exclude_frame_label_names: Blacklist of frame labels.
+        allow_external_annotations: If True, allow and automatically include annotation
+            labels that are not part of the project's official schema (e.g., labels
+            from other projects or legacy annotations). If False, these annotations
+            will be filtered out.
     """
 
     resources: Sequence['Resource']
-    resource_annotations: list[Sequence[Annotation]]
+    resource_annotations: list[Sequence['Annotation']]
     project: 'Project | None'
 
     def __init__(
@@ -89,6 +92,7 @@ def __init__(
         exclude_image_label_names: list[str] | None = None,
         include_frame_label_names: list[str] | None = None,
         exclude_frame_label_names: list[str] | None = None,
+        allow_external_annotations: bool = False,
     ):
         from datamint import Api
         # Validate mutually exclusive parameters
@@ -155,6 +159,7 @@ def __init__(
         self.exclude_image_label_names = exclude_image_label_names
         self.include_frame_label_names = include_frame_label_names
         self.exclude_frame_label_names = exclude_frame_label_names
+        self.allow_external_annotations = allow_external_annotations
 
         # Internal state
         self._logged_uint16_conversion = False
@@ -164,7 +169,7 @@ def __init__(
 
     def _extract_image_labels(
         self,
-        annotations: Sequence[Annotation],
+        annotations: Sequence['Annotation'],
     ) -> dict[str, torch.Tensor]:
         """Extract image-level label annotations.
 
@@ -263,12 +268,51 @@ def _setup_labels(self) -> None:
             self.seglabel_list, self.seglabel2code = self._process_segmentation_group(
                 worklist_schema['segmentation_group']
             )
+
+            if self.allow_external_annotations:
+                self._augment_labels_from_annotations()
         else:
             _LOGGER.info("No project provided; inferring labels from annotations.")
             self.frame_lsets, self.frame_lcodes = self._infer_labels_set(framed=True)
             self.image_lsets, self.image_lcodes = self._infer_labels_set(framed=False)
             self.seglabel_list, self.seglabel2code = self._infer_segmentation_group()
 
+    def _augment_labels_from_annotations(self) -> None:
+        """Augment project-defined label sets with identifiers found in actual annotations.
+
+        Scans resource annotations for identifiers not present in the project's
+        annotations_specs and adds them to the corresponding label/segmentation mappings.
+        """
+        inferred_frame_lsets, inferred_frame_lcodes = self._infer_labels_set(framed=True)
+        inferred_image_lsets, inferred_image_lcodes = self._infer_labels_set(framed=False)
+        inferred_seglabel_list, _ = self._infer_segmentation_group()
+
+        # Augment frame labels
+        for kind in ('multilabel', 'multiclass'):
+            existing = set(self.frame_lsets[kind])
+            new_labels = sorted([label for label in inferred_frame_lsets[kind] if label not in existing])
+            for label in new_labels:
+                _LOGGER.info(f"Allowing external frame label '{label}' not in project specs.")
+                self.frame_lsets[kind].append(label)
+            self.frame_lcodes[kind] = self.__build_label_codemap(self.frame_lsets[kind])
+
+        # Augment image labels
+        for kind in ('multilabel', 'multiclass'):
+            existing = set(self.image_lsets[kind])
+            new_labels = sorted([label for label in inferred_image_lsets[kind] if label not in existing])
+            for label in new_labels:
+                _LOGGER.info(f"Allowing external image label '{label}' not in project specs.")
+                self.image_lsets[kind].append(label)
+            self.image_lcodes[kind] = self.__build_label_codemap(self.image_lsets[kind])
+
+        # Augment segmentation labels
+        existing_segs = set(self.seglabel_list)
+        new_segs = sorted([label for label in inferred_seglabel_list if label not in existing_segs])
+        for label in new_segs:
+            _LOGGER.info(f"Allowing external segmentation label '{label}' not in project specs.")
+            self.seglabel_list.append(label)
+            self.seglabel2code[label] = len(self.seglabel_list)  # 1-based code
+
     def _setup_annotation_processor(self) -> None:
         """Initialize the annotation processor.
 
@@ -279,6 +323,7 @@ def _setup_annotation_processor(self) -> None:
             seglabel2code=self.seglabel2code,
             image_labels_set=self.image_labels_set,
             image_lcodes=self.image_lcodes,
+            allow_external_annotations=self.allow_external_annotations,
         )
 
     def _apply_annotation_filters(self) -> None:
@@ -312,11 +357,11 @@ def _filter_unannotated(self) -> None:
         self.resources = filtered_resources
         self.resource_annotations = filtered_annotations
 
-    def _filter_annotations(self, annotations: Sequence[Annotation]) -> list[Annotation]:
+    def _filter_annotations(self, annotations: Sequence['Annotation']) -> list['Annotation']:
         """Filter annotations based on include/exclude settings."""
         return [ann for ann in annotations if self._should_include_annotation(ann)]
 
-    def _should_include_annotation(self, ann: Annotation) -> bool:
+    def _should_include_annotation(self, ann: 'Annotation') -> bool:
         """Check if annotation should be included."""
         # Check annotator
         annotator = ann.created_by
@@ -331,6 +376,12 @@ def _should_include_annotation(self, ann: Annotation) -> bool:
                 return self._should_include_image_label(ann.identifier)
             else:  # frame-level
                 return self._should_include_frame_label(ann.identifier)
+        elif ann.annotation_type == 'category':
+            if not self.allow_external_annotations:
+                lsets = self.image_lsets if ann.frame_index is None else self.frame_lsets
+                valid_identifiers = {ident for ident, _ in lsets.get('multiclass', [])}
+                if ann.identifier not in valid_identifiers:
+                    return False
 
         return True
 
@@ -342,20 +393,26 @@ def _should_include_annotator(self, annotator_id: str) -> bool:
         return True
 
     def _should_include_segmentation(self, name: str) -> bool:
+        if not self.allow_external_annotations and name not in self.segmentation_labels_set:
+            return False
         if self.include_segmentation_names is not None:
             return name in self.include_segmentation_names
         if self.exclude_segmentation_names is not None:
             return name not in self.exclude_segmentation_names
         return True
 
     def _should_include_image_label(self, name: str) -> bool:
+        if not self.allow_external_annotations and name not in self.image_labels_set:
+            return False
         if self.include_image_label_names is not None:
             return name in self.include_image_label_names
         if self.exclude_image_label_names is not None:
             return name not in self.exclude_image_label_names
         return True
 
     def _should_include_frame_label(self, name: str) -> bool:
+        if not self.allow_external_annotations and name not in self.frame_labels_set:
+            return False
         if self.include_frame_label_names is not None:
             return name in self.include_frame_label_names
         if self.exclude_frame_label_names is not None:
@@ -505,7 +562,16 @@ def _preprocess_image_array(self, img: np.ndarray) -> np.ndarray:
 
     def _process_segmentations(self,
                                segmentations: dict,
-                               seg_labels: dict) -> tuple[Tensor | np.ndarray | dict, dict | None]:
+                               seg_labels: dict,
+                               output_shape: tuple | None = None) -> tuple[Tensor | np.ndarray | dict, dict | None]:
+        """
+        Process segmentations by optionally converting to semantic format and applying merge strategy.
+
+        Args:
+            segmentations: Dict of annotator_id -> segmentation array (num_instances, depth, H, W).
+            seg_labels: Dict of annotator_id -> list of label names corresponding to each instance in the segmentation array.
+            output_shape: Fallback output shape to use when outputting as semantic segmentation and no segmentations are present to infer from. Should NOT have the number of classes dimension. Example: (depth, H, W)
+        """
         # segmentations['author'] shape: (#instances, depth, H, W)
         if self.return_as_semantic_segmentation:
             sem_segs = {}
@@ -518,12 +584,22 @@ def _process_segmentations(self,
             _LOGGER.debug(
                 f'Converted to semantic segmentation. Shapes: {[segmentations[a].shape for a in segmentations]}')
             if self.semantic_seg_merge_strategy:
-                if segmentations:
+                if len(segmentations) > 0:
                     segmentations = self.annotation_processor.apply_merge_strategy(
                         segmentations,
                         strategy=self.semantic_seg_merge_strategy
                     )
                     _LOGGER.debug(f"Merged segmentation shape: {segmentations.shape}")
+                else:
+                    if output_shape is None:
+                        raise ValueError("output_shape must be provided when no segmentations are present"
+                                         " to infer shape from.")
+                    # Create empty semantic segmentation with just background class
+                    segmentations = torch.zeros((len(self.segmentation_labels_set), *output_shape),
+                                                dtype=torch.get_default_dtype())
+                    segmentations[0] = 1  # background
+                    _LOGGER.debug("No segmentations found. "
+                                  f"Created empty semantic segmentation with shape: {segmentations.shape}")
 
                 # In semantic format, we don't need `seg_labels`, as the label info is at the new dimension (axis=0) of the semantic segmentation array.
                 seg_labels = None
@@ -560,7 +636,8 @@ def __getitem__(self, index: int) -> dict[str, Any]:
                 _LOGGER.debug(
                     f"Applied albumentations transform. Image shape: {img.shape} and segs shape: {[segmentations[a].shape for a in segmentations]}")
 
-            segmentations, seg_labels = self._process_segmentations(segmentations, seg_labels)
+            segmentations, seg_labels = self._process_segmentations(segmentations, seg_labels,
+                                                                    output_shape=img.shape[1:])
 
             result['segmentations'] = segmentations
             if seg_labels:

datamint/dataset/sliced_dataset.py

@@ -417,7 +417,8 @@ def __getitem__(self, index: int) -> dict[str, Any]:
                 result['image'] = img
                 sliced_segs = aug_result['segmentations']
 
-            segmentations_processed, seg_labels_out = self._process_segmentations(sliced_segs, seg_labels)
+            segmentations_processed, seg_labels_out = self._process_segmentations(sliced_segs, seg_labels,
+                                                                                  output_shape=img.shape[1:])
             # remove temporary dummy dimension: (#instances, 1, DIM1, DIM2) -> (#instances, DIM1, DIM2)
             if isinstance(segmentations_processed, (Tensor, np.ndarray)):
                 segmentations_processed = segmentations_processed.squeeze(1)