Spaces:
No application file
No application file
| # -*- coding: utf-8 -*- | |
| import os | |
| join = os.path.join | |
| import random | |
| path_nii = '' # please complete path; two subfolders: images and labels | |
| path_video = None # or specify the path | |
| path_2d = None # or specify the path | |
| #%% split 3D nii data | |
| if path_nii is not None: | |
| img_path = join(path_nii, 'images') | |
| gt_path = join(path_nii, 'labels') | |
| gt_names = sorted(os.listdir(gt_path)) | |
| img_suffix = '_0000.nii.gz' | |
| gt_suffix = '.nii.gz' | |
| # split 20% data for validation and testing | |
| validation_path = join(path_nii, 'validation') | |
| os.makedirs(join(validation_path, 'images'), exist_ok=True) | |
| os.makedirs(join(validation_path, 'labels'), exist_ok=True) | |
| testing_path = join(path_nii, 'testing') | |
| os.makedirs(join(testing_path, 'images'), exist_ok=True) | |
| os.makedirs(join(testing_path, 'labels'), exist_ok=True) | |
| candidates = random.sample(gt_names, int(len(gt_names)*0.2)) | |
| # split half of test names for validation | |
| validation_names = random.sample(candidates, int(len(candidates)*0.5)) | |
| test_names = [name for name in candidates if name not in validation_names] | |
| # move validation and testing data to corresponding folders | |
| for name in validation_names: | |
| img_name = name.split(gt_suffix)[0] + img_suffix | |
| os.rename(join(img_path, img_name), join(validation_path, 'images', img_name)) | |
| os.rename(join(gt_path, name), join(validation_path, 'labels', name)) | |
| for name in test_names: | |
| img_name = name.split(gt_suffix)[0] + img_suffix | |
| os.rename(join(img_path, img_name), join(testing_path, 'images', img_name)) | |
| os.rename(join(gt_path, name), join(testing_path, 'labels', name)) | |
| ##% split 2D images | |
| if path_2d is not None: | |
| img_path = join(path_2d, 'images') | |
| gt_path = join(path_2d, 'labels') | |
| gt_names = sorted(os.listdir(gt_path)) | |
| img_suffix = '.png' | |
| gt_suffix = '.png' | |
| # split 20% data for validation and testing | |
| validation_path = join(path_2d, 'validation') | |
| os.makedirs(join(validation_path, 'images'), exist_ok=True) | |
| os.makedirs(join(validation_path, 'labels'), exist_ok=True) | |
| testing_path = join(path_2d, 'testing') | |
| os.makedirs(join(testing_path, 'images'), exist_ok=True) | |
| os.makedirs(join(testing_path, 'labels'), exist_ok=True) | |
| candidates = random.sample(gt_names, int(len(gt_names)*0.2)) | |
| # split half of test names for validation | |
| validation_names = random.sample(candidates, int(len(candidates)*0.5)) | |
| test_names = [name for name in candidates if name not in validation_names] | |
| # move validation and testing data to corresponding folders | |
| for name in validation_names: | |
| img_name = name.split(gt_suffix)[0] + img_suffix | |
| os.rename(join(img_path, img_name), join(validation_path, 'images', img_name)) | |
| os.rename(join(gt_path, name), join(validation_path, 'labels', name)) | |
| for name in test_names: | |
| img_name = name.split(gt_suffix)[0] + img_suffix | |
| os.rename(join(img_path, img_name), join(testing_path, 'images', img_name)) | |
| os.rename(join(gt_path, name), join(testing_path, 'labels', name)) | |
| #%% split video data | |
| if path_video is not None: | |
| img_path = join(path_video, 'images') | |
| gt_path = join(path_video, 'labels') | |
| gt_folders = sorted(os.listdir(gt_path)) | |
| # split 20% videos for validation and testing | |
| validation_path = join(path_video, 'validation') | |
| os.makedirs(join(validation_path, 'images'), exist_ok=True) | |
| os.makedirs(join(validation_path, 'labels'), exist_ok=True) | |
| testing_path = join(path_video, 'testing') | |
| os.makedirs(join(testing_path, 'images'), exist_ok=True) | |
| os.makedirs(join(testing_path, 'labels'), exist_ok=True) | |
| candidates = random.sample(gt_folders, int(len(gt_folders)*0.2)) | |
| # split half of test names for validation | |
| validation_names = random.sample(candidates, int(len(candidates)*0.5)) | |
| test_names = [name for name in candidates if name not in validation_names] | |
| # move validation and testing data to corresponding folders | |
| for name in validation_names: | |
| os.rename(join(img_path, name), join(validation_path, 'images', name)) | |
| os.rename(join(gt_path, name), join(validation_path, 'labels', name)) | |
| for name in test_names: | |
| os.rename(join(img_path, name), join(testing_path, 'images', name)) | |
| os.rename(join(gt_path, name), join(testing_path, 'labels', name)) | |