From 9cd9d80ab3d9542138422f33154ee58421a92088 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Tue, 6 Aug 2024 14:58:35 +0200 Subject: [PATCH] no annotated pixels in an image now supported in 2d --- nnunetv2/training/dataloading/base_data_loader.py | 4 ++-- nnunetv2/training/dataloading/data_loader_2d.py | 7 +++++-- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/nnunetv2/training/dataloading/base_data_loader.py b/nnunetv2/training/dataloading/base_data_loader.py index ae3006377..cab0a068e 100644 --- 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 index 8597a3b38..9f500d40d 100644 --- 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