Merge pull request #2 from fedbiomed/feature/1105-fix-medical-dataset…

…-loading-failures-from-image-names adding unit tests
fedbiomed · Nov 18, 2024 · 3bacee2 · 3bacee2
2 parents f8a49e8 + c1a925c
commit 3bacee2
Showing 1 changed file with 47 additions and 2 deletions.
diff --git a/tests/test_medical_datasets.py b/tests/test_medical_datasets.py
@@ -250,8 +250,29 @@ def _create_synthetic_dataset(self):
                 self.sample_paths.append(Path(img_path).expanduser().resolve())
                 self.sample_class.append(self.class_names.index(class_name))
 
+def _generate_image_names(add_dots_in_title, n_dots: int = 3, extension: str = '.nii.gz'):
+    # see issue 1105
+    if add_dots_in_title:
+        img_name = str(uuid4())
+        _possible_idx = list(range(len(img_name)))
+        for i in range(n_dots):
+            rand_idx = random.choice(_possible_idx)
+            img_name = img_name[:rand_idx] + '.' + img_name[rand_idx + 1:]
+            _possible_idx.remove(rand_idx)
+
+        return f'{img_name}{extension}'
+    else:
+        return f'image_{uuid4()}{extension}'
+
 
-def _create_synthetic_dataset(root: str, n_samples: int, tabular_file: str, index_col: str):
+def _create_synthetic_dataset(
+        root: str,
+        n_samples: int,
+        tabular_file: str,
+        index_col: str,
+        add_dots_in_title: bool = False,
+        extension: str = '.nii.gz'
+        ):
     """Creates synthetic dataset for test purposes
 
     Args:
@@ -280,7 +301,10 @@ def _create_synthetic_dataset(root: str, n_samples: int, tabular_file: str, inde
         for modality in modalities:
             modality_folder = os.path.join(subject_folder, modality)
             os.mkdir(modality_folder)
-            img_path = os.path.join(modality_folder, f'image_{uuid4()}.nii.gz')
+            img_path = os.path.join(
+                    modality_folder,
+                    _generate_image_names(add_dots_in_title, n_dots=6, extension=extension)
+                )
             itk.imwrite(img, img_path)
 
         # Add demographics information
@@ -793,6 +817,27 @@ def test_medical_folder_dataset_16_load_MedicalFolder(self):
             medical_folder_controller.load_MedicalFolder()
         mfd_patcher.stop()
 
+    def test_medical_folder_dataset_16_(self):
+        # test for bug #1105: failure when loading image which title contains dots
+        for extension in ('.nii', '.nii.gz'):
+            self.root2 = tempfile.mkdtemp()
+            self.tabular_file = os.path.join(self.root2, 'participants.csv')
+            self.index_col = 'FOLDER_NAME'
+
+            self.transform = {'T1': Lambda(lambda x: torch.flatten(x))}
+            self.target_transform = {'label': GaussianSmooth()}
+
+            self.n_samples = 10
+            self.batch_size = 3
+
+            print(f'Dataset folder located in: {self.root2}')
+            _create_synthetic_dataset(self.root2, self.n_samples, self.tabular_file, self.index_col, add_dots_in_title=True, extension=extension)
+            dataset = MedicalFolderDataset(self.root, tabular_file=self.tabular_file, index_col=self.index_col,)
+
+            for _subj in os.listdir(self.root2):
+                if os.path.isdir(_subj): 
+                    data = dataset.load_images(Path(self.root2).joinpath(_subj), ['T1', 'T2'])
+                    self.assertListEqual(['T1', 'T2'], list(data.keys()))
 
 class TestMedicalFolderBase(unittest.TestCase):