Accepts uncompressed nifti files
@@ -3,7 +3,6 @@ import random
import nibabel
import numpy as np
-import nibabel
import tensorflow as tf
@@ -36,14 +35,15 @@ def load_file(img_name, seg_name=None):
def list_nifti_files(img_dir, rand=False):
# TODO check images and labels consistency
train_names = [fn for fn in os.listdir(img_dir)
- if fn.endswith(".nii.gz")]
- if train_names == []:
+ if fn.lower().endswith((".nii", ".nii.gz"))]
+ if not train_names:
print('no files in {}'.format(img_dir))
raise IOError
if rand:
random.shuffle(train_names)
return train_names
+
def has_bad_inputs(args):
print('Input params:')
for arg in vars(args):
@@ -97,4 +97,3 @@ def save_segmentation(param, img_name, pred_img):
predicted_nii.set_data_dtype(np.dtype(np.float32))
nibabel.save(predicted_nii, save_name)
print('saved %s' % save_name)
-
