import os
import tables
import numpy as np
import nibabel as nib
from tqdm import tqdm
from glob import glob
from config import cfg
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 brain volumes. 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 create_table(dataset_dir, table_data_shape, save_dir, crop_coordinates, data_channels, k_fold=None):
"""
Reads and saves all brain volumes into a single table file.
Parameters
----------
dataset_dir :
The path to all brain volumes (ex: suppose we have a folder 'BraTS2019' that
contains two HGG and LGG folders each of which contains some folders so:
dataset_dir="./BraTS2019/*/*")
table_data_shape : tuple
A tuple which shows the final brain volume shape in the table
data_channels : int
Number of data channels/modalities
save_dir : str
The path to save table.
crop_coordinates : dict
k_fold : int
k-fold cross-validation
if specified, k .npy files will be saved. Each of these files shows the indexes of
brain volumes in that fold, which will be used for training the model.
Returns
-------
None
"""
all_brains_dir = glob(dataset_dir)
all_brains_dir.sort()
hdf5_file = tables.open_file(os.path.join(save_dir + 'data.hdf5'), mode='w')
filters = tables.Filters(complevel=5, complib='blosc')
data_shape = tuple([0] + list(table_data_shape) + [data_channels])
truth_shape = tuple([0] + list(table_data_shape))
affine_shape = tuple([0] + [4, 4])
data_storage = hdf5_file.create_earray(hdf5_file.root, 'data', tables.UInt16Atom(), shape=data_shape,
filters=filters, expectedrows=len(all_brains_dir))
truth_storage = hdf5_file.create_earray(hdf5_file.root, 'truth', tables.UInt8Atom(), shape=truth_shape,
filters=filters, expectedrows=len(all_brains_dir))
affine_storage = hdf5_file.create_earray(hdf5_file.root, 'affine', tables.Float32Atom(), shape=affine_shape,
filters=filters, expectedrows=len(all_brains_dir))
brain_names = []
for brain_dir in tqdm(all_brains_dir):
all_modalities, brain_affine, brain_name = read_brain(brain_dir, mode='train', **crop_coordinates)
brain = all_modalities[..., :4]
gt = all_modalities[..., -1]
# in BraTS 2017, 2018, 2019 there is no '3' label!
gt[gt==4] = 3
brain_names.append(brain_name)
data_storage.append(brain[np.newaxis,...])
truth_storage.append(gt[np.newaxis,...])
affine_storage.append(brain_affine[np.newaxis,...])
hdf5_file.create_array(hdf5_file.root, 'brain_names', obj=brain_names)
hdf5_file.close()
if k_fold:
validation_split = (1/k_fold) # this equal to 5-fold validation
all_HGG_names = [i for i in brain_names if 'HGG' in i]
all_LGG_names = [i for i in brain_names if 'LGG' in i]
np.random.seed(100)
np.random.shuffle(all_HGG_names)
np.random.shuffle(all_LGG_names)
HGG_val_size = int(validation_split * len(all_HGG_names))
LGG_val_size = int(validation_split * len(all_LGG_names))
for fold in range(k_fold):
chosen_HGG_val = all_HGG_names[fold*HGG_val_size:(fold+1)*HGG_val_size]
chosen_LGG_val = all_LGG_names[fold*LGG_val_size:(fold+1)*LGG_val_size]
chosen_HGG_train = [i for i in all_HGG_names if i not in chosen_HGG_val]
chosen_LGG_train = [i for i in all_LGG_names if i not in chosen_LGG_val]
# saving train_idx is enough
train = chosen_HGG_train + chosen_LGG_train
train_idx = [brain_names.index(i) for i in train]
train_idx.sort()
np.save(os.path.join(save_dir, 'fold{}_idx.npy'.format(fold)), train_idx)
if __name__ == '__main__':
create_table(cfg['data_dir'], cfg['table_data_shape'], cfg['save_data_dir'],
cfg['crop_coord'], cfg['data_channels'], cfg['k_fold'])
