-
Notifications
You must be signed in to change notification settings - Fork 18
/
deepseg.py
executable file
·150 lines (111 loc) · 6.95 KB
/
deepseg.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
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
import os, time, math
import nibabel as nib, numpy as np
import tensorflow as tf
import glob
from image_utils import *
def deeplearningseg(model_path, test_dir, atlas_dir):
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
# Import the computation graph and restore the variable values
saver = tf.train.import_meta_graph('{0}.meta'.format(model_path))
saver.restore(sess, '{0}'.format(model_path))
# Process each subject subdirectory
table_time = []
if os.path.exists('{0}/subjnames.txt'.format(test_dir)):
os.system('rm {0}/*.txt'.format(test_dir))
os.system('touch {0}/subjnames.txt'.format(test_dir))
for data in sorted(os.listdir(test_dir)):
print(data)
data_dir = os.path.join(test_dir, data)
if not os.path.isdir(data_dir):
print(' {0} is not a valid directory, Skip'.format(data_dir))
continue
file = open('{0}/subjnames.txt'.format(test_dir),'a')
file.write('{0}\n'.format(data))
file.close()
if os.path.exists('{0}/PHsegmentation_ED.gipl'.format(data_dir)):
os.system('rm {0}/*.gipl'.format(data_dir))
if os.path.exists('{0}/lvsa_.nii.gz'.format(data_dir)):
os.system('rm {0}/lvsa_*.nii.gz'.format(data_dir))
os.system('rm {0}/seg_*.nii.gz'.format(data_dir))
originalnii = glob.glob('{0}/*.nii'.format(data_dir))
if not originalnii:
print(' original nifit image does not exist, use lvsa.nii.gz')
originalnii = glob.glob('{0}/*.nii.gz'.format(data_dir))
imagePreprocessing(originalnii[0], data_dir, atlas_dir)
else:
print(' start image preprocessing ...')
imagePreprocessing(originalnii[0], data_dir, atlas_dir)
# Process ED and ES time frames
image_ED_name = '{0}/lvsa_{1}.nii.gz'.format(data_dir, 'ED')
image_ES_name = '{0}/lvsa_{1}.nii.gz'.format(data_dir, 'ES')
if not os.path.exists(image_ED_name) or not os.path.exists(image_ES_name):
print(' Image {0} or {1} does not exist. Skip.'.format(image_ED_name, image_ES_name))
continue
if os.path.exists('{0}/{1}'.format(data_dir, 'dofs')) or \
os.path.exists('{0}/{1}'.format(data_dir, 'segs')) or \
os.path.exists('{0}/{1}'.format(data_dir, 'tmps')) or \
os.path.exists('{0}/{1}'.format(data_dir, 'sizes')) or \
os.path.exists('{0}/{1}'.format(data_dir, 'motion')) or \
os.path.exists('{0}/{1}'.format(data_dir, 'vtks')):
os.system('rm -rf {0}/{1}'.format(data_dir, 'dofs'))
os.system('rm -rf {0}/{1}'.format(data_dir, 'segs'))
os.system('rm -rf {0}/{1}'.format(data_dir, 'tmps'))
os.system('rm -rf {0}/{1}'.format(data_dir, 'sizes'))
os.system('rm -rf {0}/{1}'.format(data_dir, 'motion'))
os.system('rm -rf {0}/{1}'.format(data_dir, 'vtks'))
os.mkdir('{0}/{1}'.format(data_dir, 'dofs'))
os.mkdir('{0}/{1}'.format(data_dir, 'segs'))
os.mkdir('{0}/{1}'.format(data_dir, 'tmps'))
os.mkdir('{0}/{1}'.format(data_dir, 'sizes'))
os.mkdir('{0}/{1}'.format(data_dir, 'motion'))
os.mkdir('{0}/{1}'.format(data_dir, 'vtks'))
else:
os.mkdir('{0}/{1}'.format(data_dir, 'dofs'))
os.mkdir('{0}/{1}'.format(data_dir, 'segs'))
os.mkdir('{0}/{1}'.format(data_dir, 'tmps'))
os.mkdir('{0}/{1}'.format(data_dir, 'sizes'))
os.mkdir('{0}/{1}'.format(data_dir, 'motion'))
os.mkdir('{0}/{1}'.format(data_dir, 'vtks'))
for fr in ['ED', 'ES']:
image_name = '{0}/lvsa_{1}.nii.gz'.format(data_dir, fr)
# Read the image
print(' Reading {} ...'.format(image_name))
nim = nib.load(image_name)
image = nim.get_data()
imageOrg = np.squeeze(image, axis=-1).astype(np.int16)
tmp = imageOrg
X, Y, Z = image.shape[:3]
print(' Segmenting {0} frame ...'.format(fr))
# print(' Segmenting {0} frame {1} ...'.format(fr, slice))
start_seg_time = time.time()
for slice in range(Z):
image = imageOrg[:,:,slice]
if image.ndim == 2:
image = np.expand_dims(image, axis=2)
# Intensity rescaling
image = rescale_intensity(image, (1, 99))
# Pad the image size to be a factor of 16 so that the downsample and upsample procedures
# in the network will result in the same image size at each resolution level.
X2, Y2 = int(math.ceil(X / 16.0)) * 16, int(math.ceil(Y / 16.0)) * 16
x_pre, y_pre = int((X2 - X) / 2), int((Y2 - Y) / 2)
x_post, y_post = (X2 - X) - x_pre, (Y2 - Y) - y_pre
image = np.pad(image, ((x_pre, x_post), (y_pre, y_post), (0, 0)), 'constant')
# Transpose the shape to NXYC
image = np.transpose(image, axes=(2, 0, 1)).astype(np.float32)
image = np.expand_dims(image, axis=-1)
# Evaluate the networ
prob, pred = sess.run(['probE:0', 'predR:0'], feed_dict={'image:0': image, 'training:0': False})
# Transpose and crop the segmentation to recover the original size
pred = np.transpose(pred, axes=(1, 2, 0))
pred = pred[x_pre:x_pre + X, y_pre:y_pre + Y]
pred = np.squeeze(pred, axis=-1).astype(np.int16)
tmp[:,:,slice] = pred
seg_time = time.time() - start_seg_time
print(' Segmentation time = {:3f}s'.format(seg_time))
table_time += [seg_time]
pred = tmp
nim2 = nib.Nifti1Image(pred, nim.affine)
nim2.header['pixdim'] = nim.header['pixdim']
nib.save(nim2, '{0}/segs/seg_lvsa_{1}.nii.gz'.format(data_dir, fr))
print('Average segmentation time = {:.3f}s per frame'.format(np.mean(table_time)))