-
Notifications
You must be signed in to change notification settings - Fork 441
/
Copy pathsplit.py
90 lines (84 loc) · 4.28 KB
/
split.py
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
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
# -*- 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))