no annotated pixels in an image now supported in 2d

MIC-DKFZ · Aug 6, 2024 · 9cd9d80 · 9cd9d80
1 parent 834b80f
commit 9cd9d80
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/nnunetv2/training/dataloading/base_data_loader.py b/nnunetv2/training/dataloading/base_data_loader.py
@@ -88,9 +88,9 @@ def get_bbox(self, data_shape: np.ndarray, force_fg: bool, class_locations: Unio
         else:
             if not force_fg and self.has_ignore:
                 selected_class = self.annotated_classes_key
-                if len(class_locations[selected_class]) == 0:
+                if class_locations is None or len(class_locations[selected_class]) == 0:
                     # no annotated pixels in this case. Not good. But we can hardly skip it here
-                    print('Warning! No annotated pixels in image!')
+                    # print('Warning! No annotated pixels in image!')
                     selected_class = None
                 # print(f'I have ignore labels and want to pick a labeled area. annotated_classes_key: {self.annotated_classes_key}')
             elif force_fg:

diff --git a/nnunetv2/training/dataloading/data_loader_2d.py b/nnunetv2/training/dataloading/data_loader_2d.py
@@ -24,7 +24,8 @@ def generate_train_batch(self):
             # select a class/region first, then a slice where this class is present, then crop to that area
             if not force_fg:
                 if self.has_ignore:
-                    selected_class_or_region = self.annotated_classes_key
+                    selected_class_or_region = self.annotated_classes_key if (
+                            len(properties['class_locations'][self.annotated_classes_key]) > 0) else None
                 else:
                     selected_class_or_region = None
             else:
@@ -41,6 +42,7 @@ def generate_train_batch(self):
 
                 selected_class_or_region = eligible_classes_or_regions[np.random.choice(len(eligible_classes_or_regions))] if \
                     len(eligible_classes_or_regions) > 0 else None
+
             if selected_class_or_region is not None:
                 selected_slice = np.random.choice(properties['class_locations'][selected_class_or_region][:, 1])
             else:
@@ -63,9 +65,10 @@ def generate_train_batch(self):
             # print(properties)
             shape = data.shape[1:]
             dim = len(shape)
-            bbox_lbs, bbox_ubs = self.get_bbox(shape, force_fg if selected_class_or_region is not None else None,
+            bbox_lbs, bbox_ubs = self.get_bbox(shape, force_fg if selected_class_or_region is not None else False,
                                                class_locations, overwrite_class=selected_class_or_region)
 
+
             # whoever wrote this knew what he was doing (hint: it was me). We first crop the data to the region of the
             # bbox that actually lies within the data. This will result in a smaller array which is then faster to pad.
             # valid_bbox is just the coord that lied within the data cube. It will be padded to match the patch size