diff --git a/nnunetv2/experiment_planning/verify_dataset_integrity.py b/nnunetv2/experiment_planning/verify_dataset_integrity.py index 5a25840f8..1359a9769 100644 --- a/nnunetv2/experiment_planning/verify_dataset_integrity.py +++ b/nnunetv2/experiment_planning/verify_dataset_integrity.py @@ -58,6 +58,13 @@ def check_cases(image_files: List[str], label_file: str, expected_num_channels: print(f'Segmentation contains NaN pixel values. You need to fix that.\nSegmentation:\n{label_file}') ret = False + # check segmentation is single-channel + if segmentation.shape[0] != 1: + print('Error: Segmentation must be a single-channel integer label map but has %d channels. ' + 'Please re-save your label file as a single-channel image with integer class indices ' + '(0, 1, 2, ...). \nSeg file: %s\n' % (segmentation.shape[0], label_file)) + ret = False + # check shapes shape_image = images.shape[1:] shape_seg = segmentation.shape[1:] diff --git a/nnunetv2/imageio/natural_image_reader_writer.py b/nnunetv2/imageio/natural_image_reader_writer.py index cbf8e614d..57abc5939 100644 --- a/nnunetv2/imageio/natural_image_reader_writer.py +++ b/nnunetv2/imageio/natural_image_reader_writer.py @@ -59,7 +59,15 @@ def read_images(self, image_fnames: Union[List[str], Tuple[str, ...]]) -> Tuple[ return np.vstack(images, dtype=np.float32, casting='unsafe'), {'spacing': (999, 1, 1)} def read_seg(self, seg_fname: str) -> Tuple[np.ndarray, dict]: - return self.read_images((seg_fname, )) + seg, props = self.read_images((seg_fname, )) + if seg.shape[0] != 1: + raise RuntimeError( + f"Segmentation file {seg_fname} has {seg.shape[0]} channels. nnU-Net expects " + f"segmentation files to be single-channel integer label maps. Please re-save your " + f"label files as grayscale (single-channel) images with integer class indices " + f"(0, 1, 2, ...)." + ) + return seg, props def write_seg(self, seg: np.ndarray, output_fname: str, properties: dict) -> None: io.imsave(output_fname, seg[0].astype(np.uint8 if np.max(seg) < 255 else np.uint16, copy=False), check_contrast=False)