From 24794f949c62f34e0b91094f60bba5a68cd099c0 Mon Sep 17 00:00:00 2001 From: Haofu Liao Date: Thu, 1 Aug 2019 23:45:03 -0400 Subject: [PATCH] Fixed a bug due to inconsistent spineweb image extension --- adn/datasets/spineweb.py | 2 +- prepare_spineweb.py | 7 +++---- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/adn/datasets/spineweb.py b/adn/datasets/spineweb.py index c973006..bcdc20f 100644 --- 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 index b4ecb33..c392edd 100644 --- 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!") \ No newline at end of file