-
Notifications
You must be signed in to change notification settings - Fork 13
/
data_prepare.py
115 lines (94 loc) · 4.12 KB
/
data_prepare.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
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
"""
Author: Benny
Date: Nov 2019
"""
import argparse
import os
import nibabel as nib
from pathlib import Path
import sys
from tqdm import tqdm
import numpy as np
def main():
for data in [x for x in os.listdir('./data/ribfrac/ribfrac-train-images-1/Part1/')]:
try:
source = nib.load('./data/ribfrac/ribfrac-train-images-1/Part1/'+data)
source = source.get_fdata()
source[source >= 200] = 1
source[source != 1] = 0
label = nib.load('./data/RibSeg/nii/'+data[:-12]+'rib-seg.nii.gz')
label = label.get_fdata()
temp = np.argwhere(source == 1)
# choice = np.random.choice(temp.shape[0], 30000, replace=False)
## downsample
# points = temp[choice, :]
label_selected_points = []
for i in temp:
label_selected_points.append(label[i[0]][i[1]][i[2]])
label_selected_points = np.array(label_selected_points)
np.save('./data/pn/data_pn/train'+data[:-13], temp)
np.save('./data/pn/label_pn/train' + data[:-13], label_selected_points)
except:
print(data,"has no label file.")
for data in [x for x in os.listdir('./data/ribfrac/ribfrac-train-images-2/Part2/')]:
try:
source = nib.load('./data/ribfrac/ribfrac-train-images-2/Part2/'+data)
source = source.get_fdata()
source[source >= 200] = 1
source[source != 1] = 0
label = nib.load('./data/RibSeg/nii/'+data[:-12]+'rib-seg.nii.gz')
label = label.get_fdata()
temp = np.argwhere(source == 1)
# choice = np.random.choice(temp.shape[0], 30000, replace=False)
# # downsample
# points = temp[choice, :]
label_selected_points = []
for i in temp:
label_selected_points.append(label[i[0]][i[1]][i[2]])
label_selected_points = np.array(label_selected_points)
np.save('./data/pn/data_pn/train'+data[:-13], temp)
np.save('./data/pn/label_pn/train' + data[:-13], label_selected_points)
except:
print(data,"has no label file.")
for data in [x for x in os.listdir('./data/ribfrac/ribfrac-val-images/')]:
try:
source = nib.load('./ribfrac/ribfrac-val-images/' + data)
source = source.get_fdata()
source[source >= 200] = 1
source[source != 1] = 0
label = nib.load('./data/RibSeg/nii/' + data[:-12] + 'rib-seg.nii.gz')
label = label.get_fdata()
temp = np.argwhere(source == 1)
# choice = np.random.choice(temp.shape[0], 30000, replace=False)
# # downsample
# points = temp[choice, :]
label_selected_points = []
for i in temp:
label_selected_points.append(label[i[0]][i[1]][i[2]])
label_selected_points = np.array(label_selected_points)
np.save('./data/pn/data_pn/val' + data[:-13], temp)
np.save('./data/pn/label_pn/val' + data[:-13], label_selected_points)
except:
print(data,"has no label file.")
for data in [x for x in os.listdir('./data/ribfrac/ribfrac-test-images/')]:
try:
source = nib.load('./data/ribfrac/ribfrac-test-images/'+data)
source = source.get_fdata()
source[source >= 200] = 1
source[source != 1] = 0
label = nib.load('./data/RibSeg/nii/'+data[:-12]+'rib-seg.nii.gz')
label = label.get_fdata()
temp = np.argwhere(source == 1)
# choice = np.random.choice(temp.shape[0], 30000, replace=False)
# # downsample
# points = temp[choice, :]
label_selected_points = []
for i in temp:
label_selected_points.append(label[i[0]][i[1]][i[2]])
temp = np.array(temp)
np.save('./data/pn/data_pn/test'+data[:-13], temp)
np.save('./data/pn/label_pn/test' + data[:-13], label_selected_points)
except:
print(data,"has no label file.")
if __name__ == '__main__':
main()