Fix id assumption

MIC-DKFZ · Oct 11, 2023 · 9a83a2e · 9a83a2e
1 parent 7e537d5
commit 9a83a2e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nnunetv2/training/dataloading/utils.py b/nnunetv2/training/dataloading/utils.py
@@ -22,7 +22,7 @@ def find_broken_image_and_labels(
     :returns: Tuple of a set containing the case ids of the broken npy images and a set of the case ids of broken npy segmentations. 
     """
     content = os.listdir(path_to_data_dir)
-    unique_ids = [c.split("_")[0] for c in content if c.endswith("_seg.npy")]
+    unique_ids = [c[:-4] for c in content if c.endswith(".npz")]
     failed_data_ids = set()
     failed_seg_ids = set()
     for unique_id in unique_ids: