-
Notifications
You must be signed in to change notification settings - Fork 1
/
check_volume.py
164 lines (138 loc) · 7.02 KB
/
check_volume.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
151
152
153
154
155
156
157
158
159
160
161
162
163
164
import os
import numpy as np
import nibabel as nib
def load_nifti(patient, nii_dir):
try:
nifti_path = os.path.join(nii_dir, f'{patient}.nii.gz')
return nib.load(nifti_path)
except:
nifti_path = os.path.join(nii_dir, f'{patient}.nii')
print('nii file has been loaded. it is recommended to compress to .nii.gz.')
return nib.load(nifti_path)
def crop_nifti(data_dir, mask_dir, patient, center):
total_height = 256
if 'NO_REG' in data_dir :
save_dir = '/data/Nifti/In/Transformed/OCTA_SRL_NO_REG_BUT_RESIZED'
else:
save_dir = f'/data/Nifti/In/Transformed/OCTA_SRL_{total_height}_V2'
nii = load_nifti(patient, data_dir)
nii_arr = np.asarray(nii.dataobj)
nii_cropped = nii_arr[:, center-int(total_height/2):center+int(total_height/2), :]
nii_cropped_img = nib.Nifti1Image(np.uint8(nii_cropped), nii.affine, nii.header)
os.makedirs(save_dir, exist_ok=True)
# nib.save(nii_cropped_img, os.path.join(save_dir, f'{patient}.nii.gz'))
print(f"{patient} has been cropped. size: {nii_cropped.shape}")
if mask_dir is not None:
save_dir_mask = f'/data/Nifti/In/Transformed/Mask_SRL_{total_height}'
# mask_nii = load_nifti(patient, mask_dir)
mask_nii = load_nifti(patient, mask_dir)
mask_nii_arr = np.asarray(mask_nii.dataobj)
mask_srl_nii_arr = np.zeros(np.shape(mask_nii_arr))
mask_srl_nii_arr[(0<mask_nii_arr) & (mask_nii_arr<=5)] = mask_nii_arr[(0<mask_nii_arr) & (mask_nii_arr<=5)]
mask_nii_cropped = np.shape((192, total_height, 192))
mask_nii_cropped = mask_srl_nii_arr[:, center-int(total_height/2):center+int(total_height/2), :]
# This is removing layer which is far away from main parts.
mask_nii_cropped = next(volume_mask_height_filtering(mask_nii_cropped))
# filtered_mask_srl_cropped = np.zeros(np.shape(mask_nii_cropped))
# filtered_mask_srl_cropped[0<nii_cropped] = mask_nii_cropped[0<nii_cropped]
mask_srl_cropped_img = nib.Nifti1Image(mask_nii_cropped, mask_nii.affine, mask_nii.header)
os.makedirs(save_dir_mask, exist_ok=True)
nib.save(mask_srl_cropped_img, os.path.join(save_dir_mask, f'{patient}.nii.gz'))
print(f"{patient} has been cropped. size: {mask_nii_cropped.shape}")
def volume_mask_height_filtering(mask_nii):
for x in range(mask_nii.shape[0]):
for z in range(mask_nii.shape[2]):
upper_list = list(mask_nii[x, :, z].nonzero()[0])
# print(upper_list)
if np.sum(mask_nii[x, :, z].nonzero()) != 0:
if len(upper_list)==1:
pass
else:
for idx, v in enumerate(upper_list[:-1]):
if upper_list[idx+1]-v > 10:
break
# print(idx)
for i in range(idx, len(upper_list)):
mask_nii[x, upper_list[i], z] = 0
yield mask_nii
def check_height(patient_id, data_dir):
# rename_file(select_dir)
patient = patient_id
patients.append(patient)
srl = load_nifti(patient_id, nii_dir=data_dir)
srl_nii = np.asarray(srl.dataobj)
if np.max(srl_nii)==0:
pass
else:
for x in range(srl_nii.shape[0]):
for y in range(srl_nii.shape[2]):
if np.sum(srl_nii[x,:,y].nonzero()) != 0:
lower = np.min(srl_nii[x,:,y].nonzero())
upper = np.max(srl_nii[x,:,y].nonzero())
center = int((lower+upper)/2)
height = upper-lower
patients_dict[patient]['up'].append(upper)
patients_dict[patient]['bottom'].append(lower)
patients_dict[patient]['height'].append(height)
patients_dict[patient]['center'].append(center)
maxH = np.max(patients_dict[patient]['height'])
maxU = np.max(patients_dict[patient]['up'])
minB = np.min(patients_dict[patient]['bottom'])
center_z_idx = int(np.median(patients_dict[patient]['center']))
need_height = int(np.abs(maxU-center_z_idx)+np.abs(center_z_idx-minB))
patients_dict[patient]['needed_height'].append(need_height)
print(f'{patient} max non-zero index : {maxU}')
print(f'{patient} min non-zero index : {minB}')
print(f'{patient} max height index : {maxH}')
print(f'{patient} center index of z-axis : {center_z_idx}' )
print(f'{patient} need {patients_dict[patient]["needed_height"][0]} height.')
mask_dir = '/data/Nifti/In/Transformed/VolMask'
# mask_dir = None
crop_nifti(data_dir, mask_dir, patient, center_z_idx)
def dict_to_CSV(save_dir, save_name):
import pandas as pd
save_name = 'volume_info.csv'
df_path = os.path.join(save_dir, save_name)
height_list = [np.max(patients_dict[p]['height']) for p in patients]
bottom_list = [np.min(patients_dict[p]['bottom']) for p in patients]
up_list = [np.max(patients_dict[p]['up']) for p in patients]
center_list = [int(np.mean(patients_dict[p]['center'])) for p in patients]
needed_height_list = [patients_dict[p]['needed_height'][0] for p in patients]
df = pd.DataFrame({'patient': patients,
'height':height_list,
'bottom':bottom_list,
'up':up_list,
'center':center_list,
'needed_height':needed_height_list})
df.to_csv(df_path, index=True, index_label='idx', mode='w')
def file_name_change(patient, before_dir, after_dir):
fov = 'fov66' if patient < 10301 else 'fov33'
before_path = os.path.join(before_dir, f'{patient}_mask_{fov}_resized.nii.gz')
after_path = os.path.join(after_dir, f'{patient}.nii.gz')
os.system(f'mv {before_path} {after_path}')
print(f'{patient} has been re-named.')
def main():
global patients_dict, data_dir, skip_list, patients
patients = []
# data_dirs = get_dir_lab()
# skip_list = [10035, 10057, 10114, 10219, 10220]
select = [patient for patient in range(10302, 10501)]
# select = [10253]
# broken = [10219, 10220]
# patients_dict = {}
patients_dict = {p : {'up':[],'bottom':[],'center':[],'height':[], 'needed_height':[]} for p in select}
data_dirs = '/data/Nifti/In/Transformed/OCTA_SRL'
crop_dir ='/data/Nifti/In/Transformed/OCTA_SRL_256'
# data_dirs = '/data/Nifti/In/Transformed/OCTA_SRL_NO_REG_BUT_RESIZED_BEFORE_ZCROP'
# crop_dir = '/data/Nifti/In/Transformed/OCTA_SRL_NO_REG_BUT_RESIZED'
os.makedirs(crop_dir, exist_ok=True)
for patient in select:
# if patient not in skip_list:
# before_dir = '/data/Nifti/In/Transformed/VolMask'
# after_dir = before_dir
# file_name_change(patient, before_dir, after_dir)
# if patient == 10012: break
check_height( patient_id=patient,
data_dir=data_dirs)
dict_to_CSV(save_dir=crop_dir, save_name='volume_info.csv')
main()