import os
import numpy as np
import nibabel as nib
from glob import glob
from tensorflow.keras.models import load_model
def read_brain(brain_dir, mode='train', x0=42, x1=194, y0=29, y1=221, z0=2, z1=146):
"""
A function that reads and crops a brain modalities (nii.gz format)
Parameters
----------
brain_dir : string
The path to a folder that contains MRI modalities of a specific brain
mode : string
'train' or 'validation' mode. The default is 'train'.
x0, x1, y0, y1, z0, z1 : int
The coordinates to crop 3D brain volume. For example, a brain volume with the
shape [x,y,z,modalites] is cropped [x0:x1, y0:y1, z0:z1, :] to have the shape
[x1-x0, y1-y0, z1-z0, modalities]. One can calculate the x0,x1,... by calculating
none zero pixels through dataset. Note that the final three shapes must be divisible
by the network downscale rate.
Returns
-------
all_modalities : array
The cropped modalities (+ gt if mode='train')
brain_affine : array
The affine matrix of the input brain volume
brain_name : str
The name of the input brain volume
"""
brain_dir = os.path.normpath(brain_dir)
flair = glob( os.path.join(brain_dir, '*_flair*.nii.gz'))
t1 = glob( os.path.join(brain_dir, '*_t1*.nii.gz'))
t1ce = glob( os.path.join(brain_dir, '*_t1ce*.nii.gz'))
t2 = glob( os.path.join(brain_dir, '*_t2*.nii.gz'))
if mode=='train':
gt = glob( os.path.join(brain_dir, '*_seg*.nii.gz'))
modalities_dir = [flair[0], t1[0], t1ce[0], t2[0], gt[0]]
elif mode=='validation':
modalities_dir = [flair[0], t1[0], t1ce[0], t2[0]]
all_modalities = []
for modality in modalities_dir:
nifti_file = nib.load(modality)
brain_numpy = np.asarray(nifti_file.dataobj)
all_modalities.append(brain_numpy)
# all modalities have the same affine, so we take one of them (the last one in this case),
# affine is just saved for preparing the predicted nii.gz file in the future.
brain_affine = nifti_file.affine
all_modalities = np.array(all_modalities)
all_modalities = np.rint(all_modalities).astype(np.int16)
all_modalities = all_modalities[:, x0:x1, y0:y1, z0:z1]
# to fit keras channel last model
all_modalities = np.transpose(all_modalities)
# tumor grade + name
brain_name = os.path.basename(os.path.split(brain_dir)[0]) + '_' + os.path.basename(brain_dir)
return all_modalities, brain_affine, brain_name
def normalize_slice(slice):
"""
Removes 1% of the top and bottom intensities and perform
normalization on the input 2D slice.
"""
b = np.percentile(slice, 99)
t = np.percentile(slice, 1)
slice = np.clip(slice, t, b)
if np.std(slice)==0:
return slice
else:
slice = (slice - np.mean(slice)) / np.std(slice)
return slice
def normalize_volume(input_volume):
"""
Perform a slice-based normalization on each modalities of input volume.
"""
normalized_slices = np.zeros_like(input_volume).astype(np.float32)
for slice_ix in range(4):
normalized_slices[slice_ix] = input_volume[slice_ix]
for mode_ix in range(input_volume.shape[1]):
normalized_slices[slice_ix][mode_ix] = normalize_slice(input_volume[slice_ix][mode_ix])
return normalized_slices
def save_predicted_results(prediction, brain_affine, view, output_dir, z_main=155, z0=2, z1=146, y_main=240, y0=29, y1=221, x_main=240, x0=42, x1=194):
"""
Save the segmented results into a .nii.gz file, so that it can be uploaded to the BraTS server.
Note that to correctly save the segmented brains, it is necessery to set x0, x1, ... correctly.
Parameters
----------
prediction : array
The predictred brain.
brain_affine : array
The affine matrix of the predicted brain volume
view : str
'axial', 'sagittal' or 'coronal'. The 'view' is needed to reconstruct output axes.
output_dir : str
The path to save .nii.gz file.
"""
prediction = np.argmax(prediction, axis=-1).astype(np.uint16)
prediction[prediction==3] = 4
if view=="axial":
prediction = np.pad(prediction, ((z0, z_main-z1), (y0, y_main-y1), (x0, x_main-x1)), 'constant')
prediction = prediction.transpose(2,1,0)
elif view=="sagital":
prediction = np.pad(prediction, ((x0, x_main-x1), (y0, y_main-y1), (z0 , z_main-z1)), 'constant')
elif view=="coronal":
prediction = np.pad(prediction, ((y0, y_main-y1), (x0, x_main-x1), (z0 , z_main-z1)), 'constant')
prediction = prediction.transpose(1,0,2)
#
prediction_ni = nib.Nifti1Image(prediction, brain_affine)
prediction_ni.to_filename(output_dir+ '.nii.gz')
if __name__ == '__main__':
val_data_dir = '/path/to/data/*'
view = 'axial'
saved_model_dir = '/path/to/a/trained/model.hdf5' #ex './save/axial_fold0/model.hdf5'
save_pred_dir = './predict/'
batch_size = 32
if not os.path.isdir(save_pred_dir):
os.mkdir(save_pred_dir)
all_brains_dir = glob(val_data_dir)
all_brains_dir.sort()
if view == 'axial':
view_axes = (0, 1, 2, 3)
elif view == 'sagittal':
view_axes = (2, 1, 0, 3)
elif view == 'coronal':
view_axes = (1, 2, 0, 3)
else:
ValueError('unknown input view => {}'.format(view))
model = load_model(saved_model_dir, compile=False)
for brain_dir in all_brains_dir:
if os.path.isdir(brain_dir):
print("Volume ID: ", os.path.basename(brain_dir))
all_modalities, brain_affine, _ = read_brain(brain_dir, mode='validation')
all_modalities = all_modalities.transpose(view_axes)
all_modalities = normalize_volume(all_modalities)
prediction = model.predict(all_modalities, batch_size=batch_size, verbose=1)
output_dir = os.path.join(save_pred_dir, os.path.basename(brain_dir))
save_predicted_results(prediction, brain_affine, view, output_dir)
