Spaces:
Runtime error
Runtime error
| 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") | |