-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathCT_util_lib.py
118 lines (93 loc) · 4.15 KB
/
CT_util_lib.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
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Jan 18 12:22:16 2023
@author: Lei Zhang
"""
#import slicer
import nibabel as nib
import os
import glob
import re
import nrrd
import numpy as np
count = 0
import SimpleITK as sitk
#This implementation is based on sitk
def nrrd2nii_itk(nrrd_fn, dest_folder, final_name):
#loadedSegmentationNode = nrrd.read(nrrd_fn)
img = sitk.ReadImage(nrrd_fn)
os.makedirs(dest_folder, exist_ok=True)
print(dest_folder + final_name)
sitk.WriteImage(img, dest_folder + final_name)
#This is low level implementation
def nrrd2nii(ref_fn, nrrd_fn, dest_folder , prefix=''):
#gt_fn_list = glob.glob(fname+'/*Nathan Segmentation.seg*.nrrd')
#print('', gt_fn_list[0])
loadedSegmentationNode = nrrd.read(nrrd_fn)
nrrd_data = loadedSegmentationNode[0]
header = loadedSegmentationNode[1]
cube_size = header["sizes"]
image_offset = header["Segmentation_ReferenceImageExtentOffset"]
result = re.finditer(r'[\s]', image_offset)
offset_list = []
xxx = 0
for m in result:
# print(m.start(0))
#print(m.)
offset_list.append(int(image_offset[xxx:m.start(0)]))
xxx = m.start(0)+1
offset_list.append(int(image_offset[xxx:]))
# print(header["sizes"])
# print(header["Segmentation_ReferenceImageExtentOffset"])
CT_data_fn = ref_fn
CT_data_orig = nib.load(CT_data_fn)
my_data = CT_data_orig.get_fdata() -CT_data_orig.get_fdata()
my_data[offset_list[0]:offset_list[0]+header["sizes"][0],offset_list[1]:offset_list[1]+header["sizes"][1], offset_list[2]:offset_list[2] + header["sizes"][2]] = nrrd_data
seg_resized_nii = nib.Nifti1Image(my_data, CT_data_orig.affine, CT_data_orig.header)
os.makedirs(dest_folder, exist_ok=True)
nrrd_fn = os.path.basename(nrrd_fn)
seg_path = dest_folder + '/'+ prefix +'_' + '.nii.gz'
nib.save(seg_resized_nii, seg_path)
# this fuction is used to convert multilabel nrrd seg file into nnunet format.
# implementated by Lei Zhang
def nrrd_analysis(ref_fn, nrrd_fn, dest_folder , prefix=''):
loadedSegmentationNode = nrrd.read(nrrd_fn)
nrrd_data = loadedSegmentationNode[0]
header = loadedSegmentationNode[1]
cube_size = header["sizes"]
image_offset = header["Segmentation_ReferenceImageExtentOffset"]
result = re.finditer(r'[\s]', image_offset)
nrrd_data_shape = nrrd_data.shape
volume_num = nrrd_data_shape[0]
print('total label before aggregation ', np.unique(nrrd_data))
label_29_total = nrrd_data[0,:,:,:]
print('total label before aggregation ', np.unique(label_29_total))
all_label_vol = nrrd_data
#aggregate from 2 to 29
for vol_idx in range(2, min(volume_num, 29)+1):
print(vol_idx, np.unique(all_label_vol[vol_idx-1,:,:,:]))
label_29_total = label_29_total + float(vol_idx) * all_label_vol[vol_idx-1,:,:,:]
print('total label after aggregation ', np.unique(label_29_total))
offset_list = []
xxx = 0
for m in result:
# print(m.start(0))
#print(m.)
offset_list.append(int(image_offset[xxx:m.start(0)]))
xxx = m.start(0)+1
offset_list.append(int(image_offset[xxx:]))
print('type of nrrd data ', type(nrrd_data), ' dimension ', nrrd_data.shape)
print('data dimension ', header["sizes"])
print('Segmentation offset ', header["Segmentation_ReferenceImageExtentOffset"])
CT_data_fn = ref_fn
CT_data_orig = nib.load(CT_data_fn)
my_data = np.zeros(CT_data_orig.get_fdata().shape, dtype=np.uint8)
#CT_data_orig.get_fdata() -CT_data_orig.get_fdata()
my_data[offset_list[0]:offset_list[0]+header["sizes"][0+1],offset_list[1]:offset_list[1]+header["sizes"][1+1], offset_list[2]:offset_list[2] + header["sizes"][2+1]] = label_29_total
print('data type ', type(my_data[0,0,0]), type(label_29_total[0,0,0]))
seg_resized_nii = nib.Nifti1Image(my_data, CT_data_orig.affine, CT_data_orig.header)
os.makedirs(dest_folder, exist_ok=True)
nrrd_fn = os.path.basename(nrrd_fn)
seg_path = dest_folder + '/'+ prefix +'.nii.gz'
nib.save(seg_resized_nii, seg_path)