Fixed a bug due to inconsistent spineweb image extension

liaohaofu · Aug 2, 2019 · 24794f9 · 24794f9
1 parent cf30a4b
commit 24794f9
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/adn/datasets/spineweb.py b/adn/datasets/spineweb.py
@@ -12,7 +12,7 @@
 
 
 class Spineweb(torch.utils.data.Dataset):
-    def __init__(self, a_dir="data/z2ai/clinical", b_dir="data/z2ai/spineweb",
+    def __init__(self, a_dir="data/spineweb/train/artifact", b_dir="data/spineweb/train/no_artifact",
         random_flip=False, a_range=(-1000.0, 2000.0), b_range=(-1000.0, 2000.0)):
         super(Spineweb, self).__init__()
 

diff --git a/prepare_spineweb.py b/prepare_spineweb.py
@@ -40,12 +40,13 @@ def make_thumbnails(images):
 
     for patient_dir in tqdm(patient_dirs):
         patient_name = path.basename(patient_dir)
-        volume_files = read_dir(patient_dir, predicate=lambda x: x.endswith("mhd"), recursive=True)
+        volume_files = read_dir(patient_dir,
+            predicate=lambda x: x.endswith("mhd") or x.endswith("nii.gz"), recursive=True)
         for volume_file in volume_files:
             volume_obj = sitk.ReadImage(volume_file)
 
             volume = sitk.GetArrayFromImage(volume_obj)
-            volume_name = path.basename(volume_file)[:-4]
+            volume_name = path.basename(volume_file).split(".")[0]
 
             thumbnails = defaultdict(list)
             index = 0
@@ -127,5 +128,3 @@ def make_thumbnails(images):
 
     shutil.rmtree(artifact_dir)
     shutil.rmtree(no_artifact_dir)
-
-    print("Done!")