YPan0's picture
Upload folder using huggingface_hub
b6deff2 verified
# -*- 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))