-
Notifications
You must be signed in to change notification settings - Fork 0
/
distort_and_register.py
57 lines (42 loc) · 2.1 KB
/
distort_and_register.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
import sys
sys.path.insert(0,'/home/gayathri/Tanvi_Temp/for_vxm/lib/python3.8/site-packages/')
import ants
import os
import numpy as np
from tqdm import tqdm
from torchvision import transforms
from utils.distortions import distort_elastic, random_elastic_transform
from utils.common_utils import read_volseg, save_section, intensity_normalize, visualize_reg
def save_Gt_Def_Reg_sections(fullname):
transform = transforms.Compose([transforms.ToTensor(),
transforms.CenterCrop(256),
])
filename = os.path.basename(fullname)
vol, _ = read_volseg(fullname)
vol = intensity_normalize(vol)
for secnum in tqdm(range(45, vol.shape[1]-45)):
section = vol[:,secnum,:]
print(f'------------Saving Section {secnum:03d}-------------')
section_tensor = transform(section)
deformed_tensor = distort_elastic(section_tensor)
deformed_section = deformed_tensor.numpy()[0]
section_mod = section_tensor.numpy()[0]
save_section(section_mod,
'../ground_truth_sections_T2/' + filename.replace('.nii.gz', f'{secnum:03d}.npy'))
save_section(deformed_section,
'../distorted_sections_T2/' + filename.replace('.nii.gz', f'{secnum:03d}.npy'))
fi_ants = ants.from_numpy(section_mod)
mi_ants = ants.from_numpy(deformed_section)
print('Sanity Check:', section.max(), section.min(),
deformed_section.max(), deformed_section.min())
mytx = ants.registration(fixed=fi_ants, moving=mi_ants, type_of_transform = 'SyN' )
moved_section = mytx['warpedmovout'].numpy()
save_section(moved_section,
'../registered_sections_T2/' + filename.replace('.nii.gz', f'{secnum:03d}.npy'))
visualize_reg(section_mod, deformed_section, moved_section,
'../visualize_sections_T2/' + filename.replace('.nii.gz', f'{secnum:03d}.png') )
def main():
fullname = '../reg_n4/IXI030-Guys-0708-T2.nii.gz'
save_Gt_Def_Reg_sections(fullname)
if __name__ == '__main__':
main()