Spaces:
Runtime error
Runtime error
File size: 1,511 Bytes
fd601de |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
import os
IMG_EXTENSIONS = [
#'.jpg', '.JPG', '.jpeg', '.JPEG',
#'.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
'.nrrd', '.nii.gz'
]
def is_image_file(filename):
return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
def make_dataset_modality(dir, accepted_modalities = ["ct"], saved_name="source_filenames.txt"):
# it works for root path of any layer:
# data_path/Task1 or Task2/pelvis or brain
# |-patient1
# |-ct.nii.gz
# |-mr.nii.gz
# |-patient2
# |-ct.nii.gz
# |-mr.nii.gz
images = []
assert os.path.isdir(dir), '%s is not a valid directory' % dir
for roots, _, files in sorted(os.walk(dir)): # os.walk digs all folders and subfolders in all layers of dir
for file in files:
if is_image_file(file) and file.split('.')[0] in accepted_modalities:
path = os.path.join(roots, file)
images.append(path)
print(f'Found {len(images)} {accepted_modalities} files in {dir} \n')
with open(saved_name,"w") as file:
for image in images:
file.write(f'{image} \n')
return images
if __name__ == '__main__':
path = "/gpfs/bwfor/work/ws/hd_qf295-foo/synthrad"
#path = "/gpfs/bwfor/work/ws/hd_qf295-foo/synthrad/Task1/pelvis"
images_list0=make_dataset_modality(path, ['ct'],"target_filenames.txt")
images_list=make_dataset_modality(path, ['cbct', 'mr'],"source_filenames.txt")
|