Update notebooks to match the new refactorings (DAT-946)

datamint/dataset/base.py

@@ -1008,6 +1008,8 @@ def _process_segmentations(self,
                         segmentations,
                         strategy=self.semantic_seg_merge_strategy
                     )
+                    if isinstance(segmentations, np.ndarray):
+                        segmentations = torch.from_numpy(segmentations).to(torch.get_default_dtype())
                     _LOGGER.debug("Merged segmentation shape: %s", segmentations.shape)
                 else:
                     if output_shape is None:

datamint/dataset/image_dataset.py

@@ -85,7 +85,8 @@ def apply_alb_transform(
         box_labels_tensor = targets.get('box_labels')
         if boxes_tensor is not None:
             alb_kwargs['bboxes'] = boxes_tensor.tolist() if isinstance(boxes_tensor, torch.Tensor) else list(boxes_tensor)
-            alb_kwargs['box_labels'] = box_labels_tensor.tolist() if isinstance(box_labels_tensor, torch.Tensor) else list(box_labels_tensor)
+            # Use 'labels': the standard albumentations label_fields key convention.
+            alb_kwargs['labels'] = box_labels_tensor.tolist() if isinstance(box_labels_tensor, torch.Tensor) else list(box_labels_tensor)
 
         aug = self.alb_transform(**alb_kwargs)
         aug_img = aug['image']
@@ -104,10 +105,10 @@ def apply_alb_transform(
                 start += count
             result['masks'] = aug_segmentations
 
-        # Reconstruct boxes
+        # Reconstruct boxes, map 'labels' back to 'box_labels' in our output dict
         if boxes_tensor is not None:
             aug_bboxes: list = list(aug.get('bboxes', []))
-            aug_box_labels: list = list(aug.get('box_labels', []))
+            aug_box_labels: list = list(aug.get('labels', []))
             if aug_bboxes:
                 result['boxes'] = torch.tensor(aug_bboxes, dtype=torch.float32)
                 result['box_labels'] = torch.tensor(aug_box_labels, dtype=torch.int64)

datamint/dataset/sliced_dataset.py

@@ -459,7 +459,7 @@ def __getitem__(self, index: int) -> dict[str, Any]:
 
         Returns dict with:
         - 'image': np.ndarray or Tensor of shape (C, H, W).
-        - 'segmentations' (if enabled): segmentation masks with depth dimension removed.
+        - 'masks' (if enabled): segmentation masks with depth dimension removed.
         - 'image_labels': dict of annotator -> label tensor.
         """
         if index >= len(self):
@@ -482,10 +482,10 @@ def __getitem__(self, index: int) -> dict[str, Any]:
 
             # Apply albumentations if present
             if self.alb_transform:
-                aug_result = self.apply_alb_transform(img, sliced_segs)
+                aug_result = self.apply_alb_transform(img, {'masks': sliced_segs})
                 img = aug_result['image']
                 result['image'] = img
-                sliced_segs = aug_result['segmentations']
+                sliced_segs = aug_result.get('masks', {})
 
             segmentations_processed, seg_labels_out = self._process_segmentations(sliced_segs, seg_labels,
                                                                                   output_shape=img.shape[1:])
@@ -497,9 +497,9 @@ def __getitem__(self, index: int) -> dict[str, Any]:
                     if isinstance(segmentations_processed[author], (Tensor, np.ndarray)):
                         segmentations_processed[author] = segmentations_processed[author].squeeze(1)
 
-            result['segmentations'] = segmentations_processed
+            result['masks'] = segmentations_processed
             if seg_labels_out:
-                result['seg_labels'] = seg_labels_out
+                result['mask_labels'] = seg_labels_out
 
         # Process image-level labels
         result['image_labels'] = self._extract_image_labels(annotations)
@@ -520,22 +520,22 @@ def __getitem__(self, index: int) -> dict[str, Any]:
     def apply_alb_transform(
         self,
         img: np.ndarray,
-        segmentations: dict[str, np.ndarray],
+        targets: dict[str, Any],
     ) -> dict[str, Any]:
         """Apply 2D albumentations transform to a single-slice image and masks.
 
-        Uses the same approach as ImageDataset: treats the data as 2D.
-
         Args:
             img: Image array of shape (C, 1, H, W) or (C, H, W).
-            segmentations: Dict of author -> mask arrays of shape (#instances, 1, H, W) or (#instances, H, W).
+            targets: Dict with optional key ``'masks'``: author -> array (#instances, 1, H, W).
 
         Returns:
-            Dict with transformed 'image' and 'segmentations'.
+            Dict with transformed ``'image'`` and ``'masks'``.
         """
         if self.alb_transform is None:
             raise ValueError("alb_transform is not set")
 
+        segmentations = targets.get('masks', {})
+
         # Squeeze depth=1 if present
         orig_dim = img.ndim
         if orig_dim == 4:
@@ -578,10 +578,10 @@ def apply_alb_transform(
         if orig_dim == 4:
             aug_img = aug_img[:, np.newaxis, :, :]
 
-        return {
-            'image': aug_img,
-            'segmentations': aug_segmentations,
-        }
+        result: dict[str, Any] = {'image': aug_img}
+        if aug_segmentations:
+            result['masks'] = aug_segmentations
+        return result
 
     def __repr__(self) -> str:
         base = super().__repr__()

datamint/lightning/trainers/lightning_modules/detection_modules/yolox_module.py

@@ -129,7 +129,7 @@ def forward(self, x: Tensor) -> Tensor:
 
     def training_step(self, batch: dict, batch_idx: int) -> Tensor:
         images: Tensor = batch['image']
-        targets = self._build_targets(batch['boxes'], batch['labels'], images.device)
+        targets = self._build_targets(batch['boxes'], batch['box_labels'], images.device)
         outputs: dict = self.model(images, targets)
         self.log('train/loss', outputs['total_loss'], prog_bar=True, on_step=True, on_epoch=True)
         self.log('train/iou_loss', outputs['iou_loss'], on_step=False, on_epoch=True)
@@ -167,7 +167,7 @@ def validation_step(self, batch: dict, batch_idx: int) -> None:
                 })
             target_list.append({
                 'boxes': batch['boxes'][i].cpu(),
-                'labels': batch['labels'][i].cpu(),
+                'labels': batch['box_labels'][i].cpu(),
             })
 
         self.map_metric.update(pred_list, target_list)

datamint/lightning/trainers/specialized/yolox.py

@@ -118,7 +118,7 @@ def _train_transform(self) -> 'BaseCompose':
 
         bbox_params = A.BboxParams(
             format='pascal_voc',
-            label_fields=['labels', 'identifiers'],
+            label_fields=['labels'],
             min_visibility=0.0,
         )
         return A.Compose([
@@ -136,7 +136,7 @@ def _eval_transform(self) -> 'BaseCompose':
 
         bbox_params = A.BboxParams(
             format='pascal_voc',
-            label_fields=['labels', 'identifiers'],
+            label_fields=['labels'],
             min_visibility=0.0,
         )
         return A.Compose([