def dilate_lesions(labels_dir, result_dir, recompute=True):
utils.mkdir(result_dir)
path_labels = utils.list_images_in_folder(labels_dir)
for path_label in path_labels:
path_result_label = os.path.join(result_dir,
os.path.basename(path_label))
if (not os.path.isfile(path_result_label)) | recompute:
label, aff, h = utils.load_volume(path_label, im_only=False)
# define lesion, WM, and LV masks
WM = (label == 2) | (label == 41)
lesion = label == 77
LV_and_lesion = (label == 4) | lesion
# morphological operations to bridge the gaps between lesions and LV
morph_struct = utils.build_binary_structure(2, len(WM.shape))
LV_and_lesion = binary_dilation(LV_and_lesion, morph_struct)
LV_and_lesion = binary_erosion(LV_and_lesion, morph_struct)
lesion = (LV_and_lesion & WM) | lesion
label[lesion] = 77
# save new label maps
utils.save_volume(label, aff, h, path_result_label)
def build_longitudinal_consensus(labels_dir_1,
labels_dir_2,
result_dir,
recompute=True):
# create result dir
utils.mkdir(result_dir)
# list all segmentations
path_labels_1 = utils.list_files(labels_dir_1)
path_labels_2 = utils.list_files(labels_dir_2)
for path_lab_1, path_lab_2 in zip(path_labels_1, path_labels_2):
# check if result is already saved
path_result = os.path.join(result_dir, os.path.basename(path_lab_1))
if (not os.path.isfile(path_result)) | recompute:
# load volumes
lab_1, aff, h = utils.load_volume(path_lab_1, im_only=False)
lab_2 = utils.load_volume(path_lab_2)
# compute and save consensus
dist_masp_1 = edit_volumes.compute_distance_map(lab_1,
crop_margin=20)
dist_masp_2 = edit_volumes.compute_distance_map(lab_2,
crop_margin=20)
consensus = (np.mean(np.stack([dist_masp_1, dist_masp_2], axis=-1),
axis=-1) > 0) * 1
utils.save_volume(consensus, aff, h, path_result)
def paste_lesions_on_buckner(lesion_dir,
buckner_dir,
result_dir,
dilate=2,
recompute=False):
path_lesions = utils.list_images_in_folder(lesion_dir)
path_buckners = utils.list_images_in_folder(buckner_dir)
utils.mkdir(result_dir)
# loop over buckner label maps
loop_info = utils.LoopInfo(len(path_buckners), 1, 'processing', True)
for idx_buckner, path_buckner in enumerate(path_buckners):
loop_info.update(idx_buckner)
buckner_name = os.path.basename(path_buckner).replace(
'_seg', '').replace('.nii.gz', '')
buckner = utils.load_volume(path_buckner)
WM = (buckner == 2) | (buckner == 7) | (buckner == 16) | (
buckner == 41) | (buckner == 46)
# loop over challenge data
for path_lesion in path_lesions:
lesion_name = os.path.basename(path_lesion).replace(
'.samseg_and_lesions.nii.gz', '')
path_result = os.path.join(
result_dir, buckner_name + '_' + lesion_name + '.nii.gz')
if (not os.path.isfile(path_result)) | recompute:
lesion = utils.load_volume(path_lesion)
assert lesion.shape == buckner.shape, 'lesions should have same shape as buckner labels'
# define lesion, WM, and LV masks
lesion = (lesion == 77) & WM
LV_and_lesion = (buckner == 4) | lesion
# morphological operations to bridge the gaps between lesions and LV
morph_struct = utils.build_binary_structure(
dilate, len(lesion.shape))
lesion = binary_dilation(LV_and_lesion, morph_struct)
lesion = binary_erosion(lesion, morph_struct)
lesion = lesion & WM
buckner_lesions = np.where(lesion, 77, buckner)
# save map
utils.save_volume(buckner_lesions, None, None, path_result)
def preprocess_asegs(aseg_dir,
lesion_gt_dir,
list_incorrect,
list_correct,
lesion_label_in_gt=77,
dilate=2,
recompute=False):
# align asegs to gt dir (cropping to same dimension)
cropped_dir = aseg_dir + '_cropped'
edit_volumes.mri_convert_images_in_dir(aseg_dir,
cropped_dir,
interpolation='nearest',
reference_dir=lesion_gt_dir,
recompute=recompute)
# correct for aseg labels
corrected_dir = cropped_dir + '_corrected'
edit_volumes.correct_labels_in_dir(cropped_dir,
list_incorrect,
list_correct,
corrected_dir,
smooth=False,
recompute=recompute)
# list gt and aseg, and create result dir
list_lesion_labels = utils.list_images_in_folder(lesion_gt_dir)
list_aseg_labels = utils.list_images_in_folder(corrected_dir)
inpainted_dir = corrected_dir + '_lesion_inpainted'
utils.mkdir(inpainted_dir)
# loop over subjects
for path_lesion_label, path_aseg_label in zip(list_lesion_labels,
list_aseg_labels):
path_result = os.path.join(inpainted_dir,
os.path.basename(path_aseg_label))
if (not os.path.isfile(path_result)) | recompute:
# paste lesion label
lesions = utils.load_volume(path_lesion_label)
aseg_label, aff, h = utils.load_volume(path_aseg_label,
im_only=False)
lesion_mask = lesions == lesion_label_in_gt
aseg_label[lesion_mask] = 77
utils.save_volume(aseg_label, aff, h, path_result)
# dilate lesion and ventricle
dilated_dir = inpainted_dir + '_dilated'
utils.mkdir(dilated_dir)
list_inpainted_aseg = utils.list_images_in_folder(inpainted_dir)
for path_aseg in list_inpainted_aseg:
path_result = os.path.join(dilated_dir, os.path.basename(path_aseg))
if (not os.path.isfile(path_result)) | recompute:
# define lesion, WM, and LV masks
aseg, aff, h = utils.load_volume(path_aseg, im_only=False)
WM = aseg == 2
LV = aseg == 4
lesion = aseg == 77
# morphological operations to bridge the gaps between lesions and LV
morph_struct = utils.build_binary_structure(
dilate, len(aseg.shape))
dilated_LV_or_lesion = binary_dilation(LV | lesion, morph_struct)
filled_LV_or_lesion = binary_erosion(dilated_LV_or_lesion,
morph_struct)
LV = LV | (filled_LV_or_lesion & WM)
aseg[LV] = 4
# save map
utils.save_volume(aseg, aff, h, path_result)
########################################################################################################
# instantiate BrainGenerator object
brain_generator = BrainGenerator(labels_dir=path_label_maps,
generation_labels=generation_labels,
output_labels=segmentation_labels,
n_neutral_labels=n_neutral_labels,
output_shape=output_shape,
output_div_by_n=output_divisible_by_n,
prior_distributions=prior_distributions,
flipping=flipping)
utils.mkdir(result_folder)
for n in range(n_examples):
# generate new image and corresponding labels
start = time.time()
im, lab = brain_generator.generate_brain()
end = time.time()
print('deformation {0:d} took {1:.01f}s'.format(n, end - start))
# save image
utils.save_volume(
np.squeeze(im), brain_generator.aff, brain_generator.header,
os.path.join(result_folder, 'SynthSeg_image_%s.nii.gz' % n))
utils.save_volume(
np.squeeze(lab), brain_generator.aff, brain_generator.header,
os.path.join(result_folder, 'SynthSeg_labels_%s.nii.gz' % n))
# Very simple script showing how to generate new images with lab2im
import os
import numpy as np
from ext.lab2im.utils import save_volume
from SynthSeg.brain_generator import BrainGenerator
# path of the input label map
path_label_map = '../data_example/brain_label_map.nii.gz'
# path where to save the generated image
result_dir = '../generated_images'
# generate an image from the label map.
# Because the image is spatially deformed, we also output the corresponding deformed label map.
brain_generator = BrainGenerator(path_label_map)
im, lab = brain_generator.generate_brain()
# save output image and label map
if not os.path.exists(os.path.join(result_dir)):
os.mkdir(result_dir)
save_volume(np.squeeze(im), brain_generator.aff, brain_generator.header,
os.path.join(result_dir, 'brain.nii.gz'))
save_volume(np.squeeze(lab), brain_generator.aff, brain_generator.header,
os.path.join(result_dir, 'labels.nii.gz'))
def prepare_hippo_training_atlases(labels_dir,
result_dir,
image_dir=None,
image_result_dir=None,
smooth=True,
crop_margin=50,
recompute=True):
"""This function prepares training label maps from CobraLab. It first crops each atlas around the right and left
hippocampi, with a margin. It then equalises the shape of these atlases by croppping them to the size of the
smallest hippocampus. Finally it realigns the obtained atlases to FS orientation axes.
:param labels_dir: path of directory with label maps to prepare
:param result_dir: path of directory where prepared atlases will be writen
:param image_dir: (optional) path of directory with images corresponding to the label maps to prepare.
This can be sued to prepare a dataset of real images for supervised training.
:param image_result_dir: (optional) path of directory where images corresponding to prepared atlases will be writen
:param smooth: (optional) whether to smooth the final cropped label maps
:param crop_margin: (optional) margin to add around hippocampi when cropping
:param recompute: (optional) whether to recompute result files even if they already exists"""
# create results dir
if not os.path.exists(result_dir):
os.mkdir(result_dir)
tmp_result_dir = os.path.join(result_dir, 'first_cropping')
if not os.path.exists(tmp_result_dir):
os.mkdir(tmp_result_dir)
if image_dir is not None:
assert image_result_dir is not None, 'image_result_dir should not be None if image_dir is specified'
if not os.path.exists(image_result_dir):
os.mkdir(image_result_dir)
tmp_image_result_dir = os.path.join(image_result_dir, 'first_cropping')
if not os.path.exists(tmp_image_result_dir):
os.mkdir(tmp_image_result_dir)
else:
tmp_image_result_dir = None
# list labels and images
labels_paths = utils.list_images_in_folder(labels_dir)
if image_dir is not None:
path_images = utils.list_images_in_folder(image_dir)
else:
path_images = [None] * len(labels_paths)
# crop all atlases around hippo
print('\ncropping around hippo')
shape_array = np.zeros((len(labels_paths)*2, 3))
for idx, (path_label, path_image) in enumerate(zip(labels_paths, path_images)):
utils.print_loop_info(idx, len(labels_paths), 1)
# crop left hippo first
path_label_first_crop_l = os.path.join(tmp_result_dir,
os.path.basename(path_label).replace('.nii', '_left.nii'))
lab, aff, h = utils.load_volume(path_label, im_only=False)
lab_l, croppping_idx, aff_l = edit_volumes.crop_volume_around_region(lab, crop_margin,
list(range(20101, 20109)), aff=aff)
if (not os.path.exists(path_label_first_crop_l)) | recompute:
utils.save_volume(lab_l, aff_l, h, path_label_first_crop_l)
else:
lab_l = utils.load_volume(path_label_first_crop_l)
if path_image is not None:
path_image_first_crop_l = os.path.join(tmp_image_result_dir,
os.path.basename(path_image).replace('.nii', '_left.nii'))
if (not os.path.exists(path_image_first_crop_l)) | recompute:
im, aff, h = utils.load_volume(path_image, im_only=False)
im, aff = edit_volumes.crop_volume_with_idx(im, croppping_idx, aff)
utils.save_volume(im, aff, h, path_image_first_crop_l)
shape_array[2*idx, :] = np.array(lab_l.shape)
# crop right hippo and flip them
path_label_first_crop_r = os.path.join(tmp_result_dir,
os.path.basename(path_label).replace('.nii', '_right_flipped.nii'))
lab, aff, h = utils.load_volume(path_label, im_only=False)
lab_r, croppping_idx, aff_r = edit_volumes.crop_volume_around_region(lab, crop_margin,
list(range(20001, 20009)), aff=aff)
if (not os.path.exists(path_label_first_crop_r)) | recompute:
lab_r = edit_volumes.flip_volume(lab_r, direction='rl', aff=aff_r)
utils.save_volume(lab_r, aff_r, h, path_label_first_crop_r)
else:
lab_r = utils.load_volume(path_label_first_crop_r)
if path_image is not None:
path_image_first_crop_r = os.path.join(tmp_image_result_dir,
os.path.basename(path_image).replace('.nii', '_right.nii'))
if (not os.path.exists(path_image_first_crop_r)) | recompute:
im, aff, h = utils.load_volume(path_image, im_only=False)
im, aff = edit_volumes.crop_volume_with_idx(im, croppping_idx, aff)
im = edit_volumes.flip_volume(im, direction='rl', aff=aff)
utils.save_volume(im, aff, h, path_image_first_crop_r)
shape_array[2*idx+1, :] = np.array(lab_r.shape)
# list croppped files
path_labels_first_cropped = utils.list_images_in_folder(tmp_result_dir)
if tmp_image_result_dir is not None:
path_images_first_cropped = utils.list_images_in_folder(tmp_image_result_dir)
else:
path_images_first_cropped = [None] * len(path_labels_first_cropped)
# crop all label maps to same size
print('\nequalising shapes')
new_shape = np.min(shape_array, axis=0).astype('int32')
for i, (path_label, path_image) in enumerate(zip(path_labels_first_cropped, path_images_first_cropped)):
utils.print_loop_info(i, len(path_labels_first_cropped), 1)
# get cropping indices
path_lab_cropped = os.path.join(result_dir, os.path.basename(path_label))
lab, aff, h = utils.load_volume(path_label, im_only=False)
lab_shape = lab.shape
min_cropping = np.array([np.maximum(int((lab_shape[i]-new_shape[i])/2), 0) for i in range(3)])
max_cropping = np.array([min_cropping[i] + new_shape[i] for i in range(3)])
# crop labels and realign on adni format
if (not os.path.exists(path_lab_cropped)) | recompute:
lab, aff = edit_volumes.crop_volume_with_idx(lab, np.concatenate([min_cropping, max_cropping]), aff)
# realign on adni format
lab = np.flip(lab, axis=2)
aff[0:3, 0:3] = np.array([[-0.6, 0, 0], [0, 0, -0.6], [0, -0.6, 0]])
utils.save_volume(lab, aff, h, path_lab_cropped)
# crop image and realign on adni format
if path_image is not None:
path_im_cropped = os.path.join(image_result_dir, os.path.basename(path_image))
if (not os.path.exists(path_im_cropped)) | recompute:
im, aff, h = utils.load_volume(path_image, im_only=False)
im, aff = edit_volumes.crop_volume_with_idx(im, np.concatenate([min_cropping, max_cropping]), aff)
im = np.flip(im, axis=2)
aff[0:3, 0:3] = np.array([[-0.6, 0, 0], [0, 0, -0.6], [0, -0.6, 0]])
im = edit_volumes.mask_volume(im, lab)
utils.save_volume(im, aff, h, path_im_cropped)
# correct all labels to left values
print('\ncorrecting labels')
list_incorrect_labels = [77, 80, 251, 252, 253, 254, 255, 29, 41, 42, 43, 44, 46, 47, 49, 50, 51, 52, 54, 58, 60,
61, 62, 63, 7012, 20001, 20002, 20004, 20005, 20006, 20007, 20008]
list_correct_labels = [2, 3, 2, 2, 2, 2, 2, 2, 2, 3, 4, 5, 7, 8, 10, 11, 12, 13, 18, 26, 28, 2, 30, 31, 20108,
20101, 20102, 20104, 20105, 20106, 20107, 20108]
edit_volumes.correct_labels_in_dir(result_dir, list_incorrect_labels, list_correct_labels, result_dir)
# smooth labels
if smooth:
print('\nsmoothing labels')
edit_volumes.smooth_labels_in_dir(result_dir, result_dir)
def preprocess_adni_hippo(path_t1,
path_t2,
path_aseg,
result_dir,
target_res,
padding_margin=85,
remove=False,
path_freesurfer='/usr/local/freesurfer/',
verbose=True,
recompute=True):
"""This function builds a T1+T2 multimodal image from the ADNI dataset.
It first rescales intensities of each channel between 0 and 255.
It then resamples the T2 image (which are 0.4*0.4*2.0 resolution) to target resolution.
The obtained T2 is then padded in all directions by the padding_margin param (typically large 85).
The T1 and aseg are then resampled like the T2 using mri_convert.
Now that the T1, T2 and asegs are aligned and at the same resolution, we crop them around the right and left hippo.
Finally, the T1 and T2 are concatenated into one single multimodal image.
:param path_t1: path input T1 (typically at 1mm isotropic)
:param path_t2: path input T2 (typically cropped around the hippo in sagittal axis, 0.4x0.4x2.0)
:param path_aseg: path input segmentation (typically at 1mm isotropic)
:param result_dir: path of directory where prepared images and labels will be writen.
:param target_res: resolution at which to resample the label maps, and the images.
Can be a number (isotropic resolution), a sequence, or a 1d numpy array.
:param padding_margin: (optional) margin to add around hippocampi when cropping
:param remove: (optional) whether to delete temporary files. Default is True.
:param path_freesurfer: (optional) path of FreeSurfer home, to use mri_convert
:param verbose: (optional) whether to print out mri_convert output when resampling images
:param recompute: (optional) whether to recompute result files even if they already exists
"""
# create results dir
if not os.path.isdir(result_dir):
os.mkdir(result_dir)
path_test_im_right = os.path.join(result_dir, 'hippo_right.nii.gz')
path_test_aseg_right = os.path.join(result_dir, 'hippo_right_aseg.nii.gz')
path_test_im_left = os.path.join(result_dir, 'hippo_left.nii.gz')
path_test_aseg_left = os.path.join(result_dir, 'hippo_left_aseg.nii.gz')
if (not os.path.isfile(path_test_im_right)) | (not os.path.isfile(path_test_aseg_right)) | \
(not os.path.isfile(path_test_im_left)) | (not os.path.isfile(path_test_aseg_left)) | recompute:
# set up FreeSurfer
os.environ['FREESURFER_HOME'] = path_freesurfer
os.system(os.path.join(path_freesurfer, 'SetUpFreeSurfer.sh'))
mri_convert = os.path.join(path_freesurfer, 'bin/mri_convert.bin')
# rescale T1
path_t1_rescaled = os.path.join(result_dir, 't1_rescaled.nii.gz')
if (not os.path.isfile(path_t1_rescaled)) | recompute:
im, aff, h = utils.load_volume(path_t1, im_only=False)
im = edit_volumes.rescale_volume(im)
utils.save_volume(im, aff, h, path_t1_rescaled)
# rescale T2
path_t2_rescaled = os.path.join(result_dir, 't2_rescaled.nii.gz')
if (not os.path.isfile(path_t2_rescaled)) | recompute:
im, aff, h = utils.load_volume(path_t2, im_only=False)
im = edit_volumes.rescale_volume(im)
utils.save_volume(im, aff, h, path_t2_rescaled)
# resample T2 to target res
path_t2_resampled = os.path.join(result_dir, 't2_rescaled_resampled.nii.gz')
if (not os.path.isfile(path_t2_resampled)) | recompute:
str_res = ' '.join([str(r) for r in utils.reformat_to_list(target_res, length=3)])
cmd = mri_convert + ' ' + path_t2_rescaled + ' ' + path_t2_resampled + ' --voxsize ' + str_res
cmd += ' -odt float'
if not verbose:
cmd += ' >/dev/null 2>&1'
_ = os.system(cmd)
# pad T2
path_t2_padded = os.path.join(result_dir, 't2_rescaled_resampled_padded.nii.gz')
if (not os.path.isfile(path_t2_padded)) | recompute:
t2, aff, h = utils.load_volume(path_t2_resampled, im_only=False)
t2_padded = np.pad(t2, padding_margin, 'constant')
aff[:3, -1] = aff[:3, -1] - (aff[:3, :3] @ (padding_margin * np.ones((3, 1)))).T
utils.save_volume(t2_padded, aff, h, path_t2_padded)
# resample T1 and aseg accordingly
path_t1_resampled = os.path.join(result_dir, 't1_rescaled_resampled.nii.gz')
if (not os.path.isfile(path_t1_resampled)) | recompute:
cmd = mri_convert + ' ' + path_t1_rescaled + ' ' + path_t1_resampled + ' -rl ' + path_t2_padded
cmd += ' -odt float'
if not verbose:
cmd += ' >/dev/null 2>&1'
_ = os.system(cmd)
path_aseg_resampled = os.path.join(result_dir, 'aseg_resampled.nii.gz')
if (not os.path.isfile(path_aseg_resampled)) | recompute:
cmd = mri_convert + ' ' + path_aseg + ' ' + path_aseg_resampled + ' -rl ' + path_t2_padded
cmd += ' -rt nearest -odt float'
if not verbose:
cmd += ' >/dev/null 2>&1'
_ = os.system(cmd)
# crop images and concatenate T1 and T2
for lab, side in zip([17, 53], ['left', 'right']):
path_test_image = os.path.join(result_dir, 'hippo_{}.nii.gz'.format(side))
path_test_aseg = os.path.join(result_dir, 'hippo_{}_aseg.nii.gz'.format(side))
if (not os.path.isfile(path_test_image)) | (not os.path.isfile(path_test_aseg)) | recompute:
aseg, aff, h = utils.load_volume(path_aseg_resampled, im_only=False)
tmp_aseg, cropping, tmp_aff = edit_volumes.crop_volume_around_region(aseg,
margin=30,
masking_labels=lab,
aff=aff)
if side == 'right':
tmp_aseg = edit_volumes.flip_volume(tmp_aseg, direction='rl', aff=tmp_aff)
utils.save_volume(tmp_aseg, tmp_aff, h, path_test_aseg)
if (not os.path.isfile(path_test_image)) | recompute:
t1 = utils.load_volume(path_t1_resampled)
t1 = edit_volumes.crop_volume_with_idx(t1, crop_idx=cropping)
t1 = edit_volumes.mask_volume(t1, tmp_aseg, dilate=6, erode=5)
t2 = utils.load_volume(path_t2_padded)
t2 = edit_volumes.crop_volume_with_idx(t2, crop_idx=cropping)
t2 = edit_volumes.mask_volume(t2, tmp_aseg, dilate=6, erode=5)
if side == 'right':
t1 = edit_volumes.flip_volume(t1, direction='rl', aff=tmp_aff)
t2 = edit_volumes.flip_volume(t2, direction='rl', aff=tmp_aff)
test_image = np.stack([t1, t2], axis=-1)
utils.save_volume(test_image, tmp_aff, h, path_test_image)
# remove unnecessary files
if remove:
list_files_to_remove = [path_t1_rescaled,
path_t2_rescaled,
path_t2_resampled,
path_t2_padded,
path_t1_resampled,
path_aseg_resampled]
for path in list_files_to_remove:
os.remove(path)
def predict(path_images,
path_model,
segmentation_label_list,
path_segmentations=None,
path_posteriors=None,
path_volumes=None,
skip_background_volume=True,
padding=None,
cropping=None,
resample=None,
aff_ref='FS',
sigma_smoothing=0,
keep_biggest_component=False,
conv_size=3,
n_levels=5,
nb_conv_per_level=2,
unet_feat_count=24,
feat_multiplier=2,
no_batch_norm=False,
activation='elu',
gt_folder=None,
evaluation_label_list=None,
verbose=True):
"""
This function uses trained models to segment images.
It is crucial that the inputs match the architecture parameters of the trained model.
:param path_images: path of the images to segment. Can be the path to a directory or the path to a single image.
:param path_model: path ot the trained model.
:param segmentation_label_list: List of labels for which to compute Dice scores. It should contain the same values
as the segmentation label list used for training the network.
Can be a sequence, a 1d numpy array, or the path to a numpy 1d array.
:param path_segmentations: (optional) path where segmentations will be writen.
Should be a dir, if path_images is a dir, and afile if path_images is a file.
Should not be None, if path_posteriors is None.
:param path_posteriors: (optional) path where posteriors will be writen.
Should be a dir, if path_images is a dir, and afile if path_images is a file.
Should not be None, if path_segmentations is None.
:param path_volumes: (optional) path of a csv file where the soft volumes of all segmented regions will be writen.
The rows of the csv file correspond to subjects, and the columns correspond to segmentation labels.
The soft volume of a structure corresponds to the sum of its predicted probability map.
:param skip_background_volume: (optional) whether to skip computing the volume of the background. This assumes the
background correspond to the first value in label list.
:param padding: (optional) crop the images to the specified shape before predicting the segmentation maps.
If padding and cropping are specified, images are padded before being cropped.
Can be an int, a sequence or a 1d numpy array.
:param cropping: (optional) crop the images to the specified shape before predicting the segmentation maps.
If padding and cropping are specified, images are padded before being cropped.
Can be an int, a sequence or a 1d numpy array.
:param resample: (optional) resample the images to the specified resolution before predicting the segmentation maps.
Can be an int, a sequence or a 1d numpy array.
:param aff_ref: (optional) type of affine matrix of the images used for training. By default this is set to the
FreeSurfer orientation ('FS'), as it was the configuration in which SynthSeg was trained. However, the new models
are now trained on data aligned with identity vox2ras matrix, so you need to change aff_ref to 'identity'.
:param sigma_smoothing: (optional) If not None, the posteriors are smoothed with a gaussian kernel of the specified
standard deviation.
:param keep_biggest_component: (optional) whether to only keep the biggest component in the predicted segmentation.
:param conv_size: (optional) size of unet's convolution masks. Default is 3.
:param n_levels: (optional) number of levels for unet. Default is 5.
:param nb_conv_per_level: (optional) number of convolution layers per level. Default is 2.
:param unet_feat_count: (optional) number of features for the first layer of the unet. Default is 24.
:param feat_multiplier: (optional) multiplicative factor for the number of feature for each new level. Default is 2.
:param no_batch_norm: (optional) whether to deactivate batch norm. Default is False.
:param activation: (optional) activation function. Can be 'elu', 'relu'.
:param gt_folder: (optional) folder containing ground truth files for evaluation.
A numpy array containing all dice scores (labels in rows, subjects in columns) will be writen either at
segmentations_dir (if not None), or posteriors_dir.
:param evaluation_label_list: (optional) if gt_folder is True you can evaluate the Dice scores on a subset of the
segmentation labels, by providing another label list here. Can be a sequence, a 1d numpy array, or the path to a
numpy 1d array. Default is the same as segmentation_label_list.
:param verbose: (optional) whether to print out info about the remaining number of cases.
"""
assert path_model, "A model file is necessary"
assert path_segmentations or path_posteriors, "output segmentation (or posteriors) is required"
# prepare output filepaths
images_to_segment, path_segmentations, path_posteriors, path_volumes = prepare_output_files(path_images,
path_segmentations,
path_posteriors,
path_volumes)
# get label and classes lists
label_list, _ = utils.get_list_labels(label_list=segmentation_label_list, FS_sort=True)
if evaluation_label_list is None:
evaluation_label_list = segmentation_label_list
# prepare volume file if needed
if path_volumes is not None:
if skip_background_volume:
csv_header = [['subject'] + [str(lab) for lab in label_list[1:]]]
else:
csv_header = [['subject'] + [str(lab) for lab in label_list]]
with open(path_volumes, 'w') as csvFile:
writer = csv.writer(csvFile)
writer.writerows(csv_header)
csvFile.close()
# perform segmentation
net = None
previous_model_input_shape = None
for idx, (path_image, path_segmentation, path_posterior) in enumerate(zip(images_to_segment,
path_segmentations,
path_posteriors)):
if verbose:
utils.print_loop_info(idx, len(images_to_segment), 10)
# preprocess image and get information
image, aff, h, im_res, n_channels, n_dims, shape, pad_shape, cropping, crop_idx = \
preprocess_image(path_image, n_levels, cropping, padding, aff_ref=aff_ref)
model_input_shape = list(image.shape[1:])
# prepare net for first image or if input's size has changed
if (idx == 0) | (previous_model_input_shape != model_input_shape):
# check for image size compatibility
if (idx != 0) & (previous_model_input_shape != model_input_shape):
print('image of different shape as previous ones, redefining network')
previous_model_input_shape = model_input_shape
# build network
net = build_model(path_model, model_input_shape, resample, im_res, n_levels, len(label_list), conv_size,
nb_conv_per_level, unet_feat_count, feat_multiplier, no_batch_norm, activation,
sigma_smoothing)
# predict posteriors
prediction_patch = net.predict(image)
# get posteriors and segmentation
seg, posteriors = postprocess(prediction_patch, cropping, pad_shape, shape, crop_idx, n_dims, label_list,
keep_biggest_component, aff, aff_ref=aff_ref)
# compute volumes
if path_volumes is not None:
if skip_background_volume:
volumes = np.sum(posteriors[..., 1:], axis=tuple(range(0, len(posteriors.shape) - 1)))
else:
volumes = np.sum(posteriors, axis=tuple(range(0, len(posteriors.shape) - 1)))
volumes = np.around(volumes * np.prod(im_res), 3)
row = [os.path.basename(path_image).replace('.nii.gz', '')] + [str(vol) for vol in volumes]
row += [np.sum(volumes[:int(len(volumes) / 2)]), np.sum(volumes[int(len(volumes) / 2):])]
with open(path_volumes, 'a') as csvFile:
writer = csv.writer(csvFile)
writer.writerow(row)
csvFile.close()
# write results to disk
if path_segmentation is not None:
utils.save_volume(seg.astype('int'), aff, h, path_segmentation)
if path_posterior is not None:
if n_channels > 1:
new_shape = list(posteriors.shape)
new_shape.insert(-1, 1)
new_shape = tuple(new_shape)
posteriors = np.reshape(posteriors, new_shape)
utils.save_volume(posteriors.astype('float'), aff, h, path_posterior)
# evaluate
if gt_folder is not None:
if path_segmentations[0] is not None:
eval_folder = os.path.dirname(path_segmentations[0])
else:
eval_folder = os.path.dirname(path_posteriors[0])
path_result_dice = os.path.join(eval_folder, 'dice.npy')
evaluate.dice_evaluation(gt_folder, eval_folder, evaluation_label_list, path_result_dice, verbose=verbose)
def predict(path_images,
path_model,
segmentation_label_list,
dist_map=False,
path_segmentations=None,
path_posteriors=None,
path_volumes=None,
segmentation_names_list=None,
padding=None,
cropping=None,
resample=None,
aff_ref='FS',
sigma_smoothing=0,
keep_biggest_component=False,
conv_size=3,
n_levels=5,
nb_conv_per_level=2,
unet_feat_count=24,
feat_multiplier=2,
activation='elu',
gt_folder=None,
evaluation_label_list=None,
compute_distances=False,
recompute=True,
verbose=True):
"""
This function uses trained models to segment images.
It is crucial that the inputs match the architecture parameters of the trained model.
:param path_images: path of the images to segment. Can be the path to a directory or the path to a single image.
:param path_model: path ot the trained model.
:param segmentation_label_list: List of labels for which to compute Dice scores. It should contain the same values
as the segmentation label list used for training the network.
Can be a sequence, a 1d numpy array, or the path to a numpy 1d array.
:param dist_map: (optional) whether the input will contain distance maps channels (between each intenisty channels)
Default is False.
:param path_segmentations: (optional) path where segmentations will be writen.
Should be a dir, if path_images is a dir, and afile if path_images is a file.
Should not be None, if path_posteriors is None.
:param path_posteriors: (optional) path where posteriors will be writen.
Should be a dir, if path_images is a dir, and afile if path_images is a file.
Should not be None, if path_segmentations is None.
:param path_volumes: (optional) path of a csv file where the soft volumes of all segmented regions will be writen.
The rows of the csv file correspond to subjects, and the columns correspond to segmentation labels.
The soft volume of a structure corresponds to the sum of its predicted probability map.
:param segmentation_names_list: (optional) List of names correponding to the names of the segmentation labels.
Only used when path_volumes is provided. Must be of the same size as segmentation_label_list. Can be given as a
list, a numpy array of strings, or the path to such a numpy array. Default is None.
:param padding: (optional) pad the images to the specified shape before predicting the segmentation maps.
Can be an int, a sequence or a 1d numpy array.
:param cropping: (optional) crop the images to the specified shape before predicting the segmentation maps.
If padding and cropping are specified, images are padded before being cropped.
Can be an int, a sequence or a 1d numpy array.
:param resample: (optional) resample the images to the specified resolution before predicting the segmentation maps.
Can be an int, a sequence or a 1d numpy array.
:param aff_ref: (optional) type of affine matrix of the images used for training. By default this is set to the
FreeSurfer orientation ('FS'), as it was the configuration in which SynthSeg was trained. However, the new models
are now trained on data aligned with identity vox2ras matrix, so you need to change aff_ref to 'identity'.
:param sigma_smoothing: (optional) If not None, the posteriors are smoothed with a gaussian kernel of the specified
standard deviation.
:param keep_biggest_component: (optional) whether to only keep the biggest component in the predicted segmentation.
:param conv_size: (optional) size of unet's convolution masks. Default is 3.
:param n_levels: (optional) number of levels for unet. Default is 5.
:param nb_conv_per_level: (optional) number of convolution layers per level. Default is 2.
:param unet_feat_count: (optional) number of features for the first layer of the unet. Default is 24.
:param feat_multiplier: (optional) multiplicative factor for the number of feature for each new level. Default is 2.
:param activation: (optional) activation function. Can be 'elu', 'relu'.
:param gt_folder: (optional) folder containing ground truth files for evaluation.
A numpy array containing all dice scores (labels in rows, subjects in columns) will be writen either at
segmentations_dir (if not None), or posteriors_dir.
:param evaluation_label_list: (optional) if gt_folder is True you can evaluate the Dice scores on a subset of the
segmentation labels, by providing another label list here. Can be a sequence, a 1d numpy array, or the path to a
numpy 1d array. Default is the same as segmentation_label_list.
:param recompute: (optional) whether to recompute segmentations that were already computed. This also applies to
Dice scores, if gt_folder is not None. Default is True.
:param verbose: (optional) whether to print out info about the remaining number of cases.
"""
# prepare output filepaths
images_to_segment, path_segmentations, path_posteriors, path_volumes, compute = \
prepare_output_files(path_images, path_segmentations, path_posteriors, path_volumes, recompute)
# get label and classes lists
label_list, n_neutral_labels = utils.get_list_labels(label_list=segmentation_label_list, FS_sort=True)
if evaluation_label_list is None:
evaluation_label_list = segmentation_label_list
# prepare volume file if needed
if path_volumes is not None:
if segmentation_names_list is not None:
csv_header = [[''] + utils.reformat_to_list(segmentation_names_list, load_as_numpy=True)]
csv_header += [[''] + [str(lab) for lab in label_list[1:]]]
else:
csv_header = [['subjects'] + [str(lab) for lab in label_list[1:]]]
with open(path_volumes, 'w') as csvFile:
writer = csv.writer(csvFile)
writer.writerows(csv_header)
csvFile.close()
# perform segmentation
net = None
previous_model_input_shape = None
loop_info = utils.LoopInfo(len(images_to_segment), 10, 'predicting', True)
for idx, (path_image, path_segmentation, path_posterior, tmp_compute) in enumerate(zip(images_to_segment,
path_segmentations,
path_posteriors,
compute)):
# compute segmentation only if needed
if tmp_compute:
# preprocess image and get information
image, aff, h, im_res, n_channels, n_dims, shape, pad_shape, crop_idx = \
preprocess_image(path_image, n_levels, cropping, padding, aff_ref=aff_ref, dist_map=dist_map)
model_input_shape = list(image.shape[1:])
# prepare net for first image or if input's size has changed
if (net is None) | (previous_model_input_shape != model_input_shape):
# check for image size compatibility
if (net is not None) & (previous_model_input_shape != model_input_shape) & verbose:
print('image of different shape as previous ones, redefining network')
previous_model_input_shape = model_input_shape
# build network
net = build_model(path_model, model_input_shape, resample, im_res, n_levels, len(label_list), conv_size,
nb_conv_per_level, unet_feat_count, feat_multiplier, activation, sigma_smoothing)
if verbose:
loop_info.update(idx)
# predict posteriors
prediction_patch = net.predict(image)
# get posteriors and segmentation
seg, posteriors = postprocess(prediction_patch, pad_shape, shape, crop_idx, n_dims, label_list,
keep_biggest_component, aff, aff_ref=aff_ref,
keep_biggest_of_each_group=keep_biggest_component,
n_neutral_labels=n_neutral_labels)
# write results to disk
if path_segmentation is not None:
utils.save_volume(seg.astype('int'), aff, h, path_segmentation)
if path_posterior is not None:
if n_channels > 1:
posteriors = utils.add_axis(posteriors, axis=[0, -1])
utils.save_volume(posteriors.astype('float'), aff, h, path_posterior)
else:
if path_volumes is not None:
posteriors, _, _, _, _, _, im_res = utils.get_volume_info(path_posterior, True, aff_ref=np.eye(4))
else:
posteriors = im_res = None
# compute volumes
if path_volumes is not None:
volumes = np.sum(posteriors[..., 1:], axis=tuple(range(0, len(posteriors.shape) - 1)))
volumes = np.around(volumes * np.prod(im_res), 3)
row = [os.path.basename(path_image).replace('.nii.gz', '')] + [str(vol) for vol in volumes]
with open(path_volumes, 'a') as csvFile:
writer = csv.writer(csvFile)
writer.writerow(row)
csvFile.close()
# evaluate
if gt_folder is not None:
# find path evaluation folder
path_first_result = path_segmentations[0] if (path_segmentations[0] is not None) else path_posteriors[0]
eval_folder = os.path.dirname(path_first_result)
# compute evaluation metrics
evaluate.dice_evaluation(gt_folder,
eval_folder,
evaluation_label_list,
compute_distances=compute_distances,
compute_score_whole_structure=False,
path_dice=os.path.join(eval_folder, 'dice.npy'),
path_hausdorff=os.path.join(eval_folder, 'hausdorff.npy'),
path_mean_distance=os.path.join(eval_folder, 'mean_distance.npy'),
recompute=recompute,
verbose=verbose)
if args['ct']:
im[im < 0] = 0
im[im > 80] = 80
im, aff = edit_volumes.resample_volume(im, aff, [1.0, 1.0, 1.0])
aff_ref = np.eye(4)
im, aff2 = edit_volumes.align_volume_to_ref(im,
aff,
aff_ref=aff_ref,
return_aff=True,
n_dims=3)
im = im - np.min(im)
im = im / np.max(im)
I = im[np.newaxis, ..., np.newaxis]
W = (np.ceil(np.array(I.shape[1:-1]) / 32.0) * 32).astype('int')
idx = np.floor((W - I.shape[1:-1]) / 2).astype('int')
S = np.zeros([1, *W, 1])
S[0, idx[0]:idx[0] + I.shape[1], idx[1]:idx[1] + I.shape[2],
idx[2]:idx[2] + I.shape[3], :] = I
output = unet_model.predict(S)
pred = np.squeeze(output)
pred = 255 * pred
pred[pred < 0] = 0
pred[pred > 128] = 128
pred = pred[idx[0]:idx[0] + I.shape[1], idx[1]:idx[1] + I.shape[2],
idx[2]:idx[2] + I.shape[3]]
utils.save_volume(pred, aff2, None, path_prediction)
print(' ')
print('All done!')
print(' ')
def predict(path_images,
path_model,
segmentation_label_list,
path_segmentations=None,
path_posteriors=None,
path_volumes=None,
voxel_volume=1.,
skip_background_volume=True,
padding=None,
cropping=None,
resample=None,
sigma_smoothing=0,
keep_biggest_component=False,
conv_size=3,
n_levels=5,
nb_conv_per_level=2,
unet_feat_count=24,
feat_multiplier=2,
no_batch_norm=False,
gt_folder=None):
"""
This function uses trained models to segment images.
It is crucial that the inputs match the architecture parameters of the trained model.
:param path_images: path of the images to segment. Can be the path to a directory or the path to a single image.
:param path_model: path ot the trained model.
:param segmentation_label_list: List of labels for which to compute Dice scores. It should contain the same values
as the segmentation label list used for training the network.
Can be a sequence, a 1d numpy array, or the path to a numpy 1d array.
:param path_segmentations: (optional) path where segmentations will be writen.
Should be a dir, if path_images is a dir, and afile if path_images is a file.
Should not be None, if path_posteriors is None.
:param path_posteriors: (optional) path where posteriors will be writen.
Should be a dir, if path_images is a dir, and afile if path_images is a file.
Should not be None, if path_segmentations is None.
:param path_volumes: (optional) path of a csv file where the soft volumes of all segmented regions will be writen.
The rows of the csv file correspond to subjects, and the columns correspond to segmentation labels.
The soft volume of a structure corresponds to the sum of its predicted probability map.
:param voxel_volume: (optional) volume of voxel. Default is 1 (i.e. returned volumes are voxel counts).
:param skip_background_volume: (optional) whether to skip computing the volume of the background. This assumes the
background correspond to the first value in label list.
:param padding: (optional) crop the images to the specified shape before predicting the segmentation maps.
If padding and cropping are specified, images are padded before being cropped.
Can be an int, a sequence or a 1d numpy array.
:param cropping: (optional) crop the images to the specified shape before predicting the segmentation maps.
If padding and cropping are specified, images are padded before being cropped.
Can be an int, a sequence or a 1d numpy array.
:param resample: (optional) resample the images to the specified resolution before predicting the segmentation maps.
Can be an int, a sequence or a 1d numpy array.
:param sigma_smoothing: (optional) If not None, the posteriors are smoothed with a gaussian kernel of the specified
standard deviation.
:param keep_biggest_component: (optional) whether to only keep the biggest component in the predicted segmentation.
:param conv_size: (optional) size of unet's convolution masks. Default is 3.
:param n_levels: (optional) number of levels for unet. Default is 5.
:param nb_conv_per_level: (optional) number of convolution layers per level. Default is 2.
:param unet_feat_count: (optional) number of features for the first layer of the unet. Default is 24.
:param feat_multiplier: (optional) multiplicative factor for the number of feature for each new level. Default is 2.
:param no_batch_norm: (optional) whether to deactivate batch norm. Default is False.
:param gt_folder: (optional) folder containing ground truth files for evaluation.
A numpy array containing all dice scores (labels in rows, subjects in columns) will be writen either at
segmentations_dir (if not None), or posteriors_dir.
"""
assert path_model, "A model file is necessary"
assert path_segmentations or path_posteriors, "output segmentation (or posteriors) is required"
# prepare output filepaths
images_to_segment, path_segmentations, path_posteriors, path_volumes = prepare_output_files(
path_images, path_segmentations, path_posteriors, path_volumes)
# get label and classes lists
label_list, _ = utils.get_list_labels(label_list=segmentation_label_list,
FS_sort=True)
# prepare volume file if needed
if path_volumes is not None:
if skip_background_volume:
csv_header = [['subject'] + [str(lab) for lab in label_list[1:]]]
else:
csv_header = [['subject'] + [str(lab) for lab in label_list]]
with open(path_volumes, 'w') as csvFile:
writer = csv.writer(csvFile)
writer.writerows(csv_header)
csvFile.close()
# perform segmentation
net = None
previous_model_input_shape = None
for idx, (im_path, seg_path, posteriors_path) in enumerate(
zip(images_to_segment, path_segmentations, path_posteriors)):
utils.print_loop_info(idx, len(images_to_segment), 10)
# preprocess image and get information
image, aff, h, n_channels, n_dims, shape, pad_shape, cropping, crop_idx = preprocess_image(
im_path, n_levels, cropping, padding)
model_input_shape = image.shape[1:]
# prepare net for first image or if input's size has changed
if (idx == 0) | (previous_model_input_shape != model_input_shape):
# check for image size compatibility
if (idx != 0) & (previous_model_input_shape != model_input_shape):
print(
'image of different shape as previous ones, redefining network'
)
previous_model_input_shape = model_input_shape
net = None
if resample is not None:
net, resample_shape = preprocessing_model(
resample, model_input_shape, h, n_channels, n_dims,
n_levels)
else:
resample_shape = previous_model_input_shape
net = prepare_unet(resample_shape,
len(label_list),
conv_size,
n_levels,
nb_conv_per_level,
unet_feat_count,
feat_multiplier,
no_batch_norm,
path_model,
input_model=net)
if (resample is not None) | (sigma_smoothing != 0):
net = postprocessing_model(net, model_input_shape, resample,
sigma_smoothing, n_dims)
# predict posteriors
prediction_patch = net.predict(image)
# get posteriors and segmentation
seg, posteriors = postprocess(prediction_patch, cropping, pad_shape,
shape, crop_idx, n_dims, label_list,
keep_biggest_component)
# compute volumes
if path_volumes is not None:
if skip_background_volume:
volumes = np.around(
np.sum(posteriors[..., 1:],
axis=tuple(range(0,
len(posteriors.shape) - 1))), 3)
else:
volumes = np.around(
np.sum(posteriors,
axis=tuple(range(0,
len(posteriors.shape) - 1))), 3)
volumes = voxel_volume * volumes
row = [os.path.basename(im_path)] + [str(vol) for vol in volumes]
with open(path_volumes, 'a') as csvFile:
writer = csv.writer(csvFile)
writer.writerow(row)
csvFile.close()
# write results to disk
if seg_path is not None:
utils.save_volume(seg.astype('int'), aff, h, seg_path)
if posteriors_path is not None:
if n_channels > 1:
new_shape = list(posteriors.shape)
new_shape.insert(-1, 1)
new_shape = tuple(new_shape)
posteriors = np.reshape(posteriors, new_shape)
utils.save_volume(posteriors.astype('float'), aff, h,
posteriors_path)
# evaluate
if gt_folder is not None:
if path_segmentations[0] is not None:
eval_folder = os.path.dirname(path_segmentations[0])
else:
eval_folder = os.path.dirname(path_posteriors[0])
path_result_dice = os.path.join(eval_folder, 'dice.npy')
evaluate.dice_evaluation(gt_folder, eval_folder,
segmentation_label_list, path_result_dice)
If you use this code, please cite one of the SynthSeg papers:
https://github.com/BBillot/SynthSeg/blob/master/bibtex.bib
Copyright 2020 Benjamin Billot
Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in
compliance with the License. You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software distributed under the License is
distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied. See the License for the specific language governing permissions and limitations under the
License.
"""
from ext.lab2im import utils
from SynthSeg.brain_generator import BrainGenerator
# generate an image from the label map.
brain_generator = BrainGenerator(
'../../data/training_label_maps/training_seg_01.nii.gz')
im, lab = brain_generator.generate_brain()
# save output image and label map
utils.save_volume(im, brain_generator.aff, brain_generator.header,
'./generated_examples/image_default.nii.gz')
utils.save_volume(lab, brain_generator.aff, brain_generator.header,
'./generated_examples/labels_default.nii.gz')
def preprocess_image(im_path,
n_levels,
target_res,
crop=None,
padding=None,
flip=False,
path_resample=None):
# read image and corresponding info
im, _, aff, n_dims, n_channels, header, im_res = utils.get_volume_info(
im_path, True)
# resample image if necessary
if target_res is not None:
target_res = np.squeeze(
utils.reformat_to_n_channels_array(target_res, n_dims))
if np.any((im_res > target_res + 0.05) | (im_res < target_res - 0.05)):
im_res = target_res
im, aff = edit_volumes.resample_volume(im, aff, im_res)
if path_resample is not None:
utils.save_volume(im, aff, header, path_resample)
# align image
im = edit_volumes.align_volume_to_ref(im,
aff,
aff_ref=np.eye(4),
n_dims=n_dims)
shape = list(im.shape)
# pad image if specified
if padding:
im = edit_volumes.pad_volume(im, padding_shape=padding)
pad_shape = im.shape[:n_dims]
else:
pad_shape = shape
# check that patch_shape or im_shape are divisible by 2**n_levels
if crop is not None:
crop = utils.reformat_to_list(crop, length=n_dims, dtype='int')
if not all([pad_shape[i] >= crop[i] for i in range(len(pad_shape))]):
crop = [min(pad_shape[i], crop[i]) for i in range(n_dims)]
if not all([size % (2**n_levels) == 0 for size in crop]):
crop = [
utils.find_closest_number_divisible_by_m(size, 2**n_levels)
for size in crop
]
else:
if not all([size % (2**n_levels) == 0 for size in pad_shape]):
crop = [
utils.find_closest_number_divisible_by_m(size, 2**n_levels)
for size in pad_shape
]
# crop image if necessary
if crop is not None:
im, crop_idx = edit_volumes.crop_volume(im,
cropping_shape=crop,
return_crop_idx=True)
else:
crop_idx = None
# normalise image
if n_channels == 1:
im = edit_volumes.rescale_volume(im,
new_min=0.,
new_max=1.,
min_percentile=0.5,
max_percentile=99.5)
else:
for i in range(im.shape[-1]):
im[..., i] = edit_volumes.rescale_volume(im[..., i],
new_min=0.,
new_max=1.,
min_percentile=0.5,
max_percentile=99.5)
# flip image along right/left axis
if flip & (n_dims > 2):
im_flipped = edit_volumes.flip_volume(im,
direction='rl',
aff=np.eye(4))
im_flipped = utils.add_axis(
im_flipped) if n_channels > 1 else utils.add_axis(im_flipped,
axis=[0, -1])
else:
im_flipped = None
# add batch and channel axes
im = utils.add_axis(im) if n_channels > 1 else utils.add_axis(im,
axis=[0, -1])
return im, aff, header, im_res, n_channels, n_dims, shape, pad_shape, crop_idx, im_flipped
def predict(path_images,
path_segmentations,
path_model,
segmentation_labels,
n_neutral_labels=None,
path_posteriors=None,
path_resampled=None,
path_volumes=None,
segmentation_label_names=None,
padding=None,
cropping=None,
target_res=1.,
gradients=False,
flip=True,
topology_classes=None,
sigma_smoothing=0.5,
keep_biggest_component=True,
conv_size=3,
n_levels=5,
nb_conv_per_level=2,
unet_feat_count=24,
feat_multiplier=2,
activation='elu',
gt_folder=None,
evaluation_labels=None,
mask_folder=None,
list_incorrect_labels=None,
list_correct_labels=None,
compute_distances=False,
recompute=True,
verbose=True):
"""
This function uses trained models to segment images.
It is crucial that the inputs match the architecture parameters of the trained model.
:param path_images: path of the images to segment. Can be the path to a directory or the path to a single image.
:param path_segmentations: path where segmentations will be writen.
Should be a dir, if path_images is a dir, and a file if path_images is a file.
:param path_model: path ot the trained model.
:param segmentation_labels: List of labels for which to compute Dice scores. It should be the same list as the
segmentation_labels used in training.
:param n_neutral_labels: (optional) if the label maps contain some right/left specific labels and if test-time
flipping is applied (see parameter 'flip'), please provide the number of non-sided labels (including background).
It should be the same value as for training. Default is None.
:param path_posteriors: (optional) path where posteriors will be writen.
Should be a dir, if path_images is a dir, and a file if path_images is a file.
:param path_resampled: (optional) path where images resampled to 1mm isotropic will be writen.
We emphasise that images are resampled as soon as the resolution in one of the axes is not in the range [0.9; 1.1].
Should be a dir, if path_images is a dir, and a file if path_images is a file. Default is None, where resampled
images are not saved.
:param path_volumes: (optional) path of a csv file where the soft volumes of all segmented regions will be writen.
The rows of the csv file correspond to subjects, and the columns correspond to segmentation labels.
The soft volume of a structure corresponds to the sum of its predicted probability map.
:param segmentation_label_names: (optional) List of names correponding to the names of the segmentation labels.
Only used when path_volumes is provided. Must be of the same size as segmentation_labels. Can be given as a
list, a numpy array of strings, or the path to such a numpy array. Default is None.
:param padding: (optional) pad the images to the specified shape before predicting the segmentation maps.
Can be an int, a sequence or a 1d numpy array.
:param cropping: (optional) crop the images to the specified shape before predicting the segmentation maps.
If padding and cropping are specified, images are padded before being cropped.
Can be an int, a sequence or a 1d numpy array.
:param target_res: (optional) target resolution at which the network operates (and thus resolution of the output
segmentations). This must match the resolution of the training data ! target_res is used to automatically resampled
the images with resolutions outside [target_res-0.05, target_res+0.05].
Can be a sequence, a 1d numpy array. Set to None to disable the automatic resampling. Default is 1mm.
:param flip: (optional) whether to perform test-time augmentation, where the input image is segmented along with
a right/left flipped version on it. If set to True (default), be careful because this requires more memory.
:param topology_classes: List of classes corresponding to all segmentation labels, in order to group them into
classes, for each of which we will operate a smooth version of biggest connected component.
Can be a sequence, a 1d numpy array, or the path to a numpy 1d array in the same order as segmentation_labels.
Default is None, where no topological analysis is performed.
:param sigma_smoothing: (optional) If not None, the posteriors are smoothed with a gaussian kernel of the specified
standard deviation.
:param keep_biggest_component: (optional) whether to only keep the biggest component in the predicted segmentation.
This is applied independently of topology_classes, and it is applied to the whole segmentation
:param conv_size: (optional) size of unet's convolution masks. Default is 3.
:param n_levels: (optional) number of levels for unet. Default is 5.
:param nb_conv_per_level: (optional) number of convolution layers per level. Default is 2.
:param unet_feat_count: (optional) number of features for the first layer of the unet. Default is 24.
:param feat_multiplier: (optional) multiplicative factor for the number of feature for each new level. Default is 2.
:param activation: (optional) activation function. Can be 'elu', 'relu'.
:param gt_folder: (optional) path of the ground truth label maps corresponding to the input images. Should be a dir,
if path_images is a dir, or a file if path_images is a file.
Providing a gt_folder will trigger a Dice evaluation, where scores will be writen along with the path_segmentations.
Specifically, the scores are contained in a numpy array, where labels are in rows, and subjects in columns.
:param evaluation_labels: (optional) if gt_folder is True you can evaluate the Dice scores on a subset of the
segmentation labels, by providing another label list here. Can be a sequence, a 1d numpy array, or the path to a
numpy 1d array. Default is np.unique(segmentation_labels).
:param mask_folder: (optional) path of masks that will be used to mask out some parts of the obtained segmentations
during the evaluation. Default is None, where nothing is masked.
:param list_incorrect_labels: (optional) this option enables to replace some label values in the obtained
segmentations by other label values. Can be a list, a 1d numpy array, or the path to such an array.
:param list_correct_labels: (optional) list of values to correct the labels specified in list_incorrect_labels.
Correct values must have the same order as their corresponding value in list_incorrect_labels.
:param compute_distances: (optional) whether to add Hausdorff and mean surface distance evaluations to the default
Dice evaluation. Default is True.
:param recompute: (optional) whether to recompute segmentations that were already computed. This also applies to
Dice scores, if gt_folder is not None. Default is True.
:param verbose: (optional) whether to print out info about the remaining number of cases.
"""
# prepare input/output filepaths
path_images, path_segmentations, path_posteriors, path_resampled, path_volumes, compute = \
prepare_output_files(path_images, path_segmentations, path_posteriors, path_resampled, path_volumes, recompute)
# get label list
segmentation_labels, _ = utils.get_list_labels(
label_list=segmentation_labels)
n_labels = len(segmentation_labels)
# get unique label values, and build correspondance table between contralateral structures if necessary
if (n_neutral_labels is not None) & flip:
n_sided_labels = int((n_labels - n_neutral_labels) / 2)
lr_corresp = np.stack([
segmentation_labels[n_neutral_labels:n_neutral_labels +
n_sided_labels],
segmentation_labels[n_neutral_labels + n_sided_labels:]
])
segmentation_labels, indices = np.unique(segmentation_labels,
return_index=True)
lr_corresp_unique, lr_corresp_indices = np.unique(lr_corresp[0, :],
return_index=True)
lr_corresp_unique = np.stack(
[lr_corresp_unique, lr_corresp[1, lr_corresp_indices]])
lr_corresp_unique = lr_corresp_unique[:, 1:] if not np.all(
lr_corresp_unique[:, 0]) else lr_corresp_unique
lr_indices = np.zeros_like(lr_corresp_unique)
for i in range(lr_corresp_unique.shape[0]):
for j, lab in enumerate(lr_corresp_unique[i]):
lr_indices[i, j] = np.where(segmentation_labels == lab)[0]
else:
segmentation_labels, indices = np.unique(segmentation_labels,
return_index=True)
lr_indices = None
# prepare topology classes
if topology_classes is not None:
topology_classes = utils.load_array_if_path(
topology_classes, load_as_numpy=True)[indices]
# prepare volume file if needed
if path_volumes is not None:
if segmentation_label_names is not None:
segmentation_label_names = utils.load_array_if_path(
segmentation_label_names)[indices]
csv_header = [[''] + segmentation_label_names[1:].tolist()]
csv_header += [[''] +
[str(lab) for lab in segmentation_labels[1:]]]
else:
csv_header = [['subjects'] +
[str(lab) for lab in segmentation_labels[1:]]]
with open(path_volumes, 'w') as csvFile:
writer = csv.writer(csvFile)
writer.writerows(csv_header)
csvFile.close()
# build network
_, _, n_dims, n_channels, _, _ = utils.get_volume_info(path_images[0])
model_input_shape = [None] * n_dims + [n_channels]
net = build_model(path_model, model_input_shape, n_levels,
len(segmentation_labels), conv_size, nb_conv_per_level,
unet_feat_count, feat_multiplier, activation,
sigma_smoothing, gradients)
# perform segmentation
loop_info = utils.LoopInfo(len(path_images), 10, 'predicting', True)
for idx, (path_image, path_segmentation, path_posterior, path_resample, tmp_compute) in \
enumerate(zip(path_images, path_segmentations, path_posteriors, path_resampled, compute)):
# compute segmentation only if needed
if tmp_compute:
if verbose:
loop_info.update(idx)
# preprocessing
image, aff, h, im_res, _, _, shape, pad_shape, crop_idx, im_flipped = \
preprocess_image(path_image, n_levels, target_res, cropping, padding, flip, path_resample)
# prediction
prediction_patch = net.predict(image)
prediction_patch_flip = net.predict(im_flipped) if flip else None
# postprocessing
seg, posteriors = postprocess(
prediction_patch,
pad_shape,
shape,
crop_idx,
n_dims,
segmentation_labels,
lr_indices,
keep_biggest_component,
aff,
topology_classes=topology_classes,
post_patch_flip=prediction_patch_flip)
# write results to disk
if path_segmentation is not None:
utils.save_volume(seg,
aff,
h,
path_segmentation,
dtype='int32')
if path_posterior is not None:
if n_channels > 1:
posteriors = utils.add_axis(posteriors, axis=[0, -1])
utils.save_volume(posteriors,
aff,
h,
path_posterior,
dtype='float32')
else:
if path_volumes is not None:
posteriors, _, _, _, _, _, im_res = utils.get_volume_info(
path_posterior, True, aff_ref=np.eye(4))
else:
posteriors = im_res = None
# compute volumes
if path_volumes is not None:
volumes = np.sum(posteriors[..., 1:],
axis=tuple(range(0,
len(posteriors.shape) - 1)))
volumes = np.around(volumes * np.prod(im_res), 3)
row = [os.path.basename(path_image).replace('.nii.gz', '')
] + [str(vol) for vol in volumes]
with open(path_volumes, 'a') as csvFile:
writer = csv.writer(csvFile)
writer.writerow(row)
csvFile.close()
# evaluate
if gt_folder is not None:
# find path where segmentations are saved evaluation folder, and get labels on which to evaluate
eval_folder = os.path.dirname(path_segmentations[0])
if evaluation_labels is None:
evaluation_labels = segmentation_labels
# set path of result arrays for surface distance if necessary
if compute_distances:
path_hausdorff = os.path.join(eval_folder, 'hausdorff.npy')
path_hausdorff_99 = os.path.join(eval_folder, 'hausdorff_99.npy')
path_hausdorff_95 = os.path.join(eval_folder, 'hausdorff_95.npy')
path_mean_distance = os.path.join(eval_folder, 'mean_distance.npy')
else:
path_hausdorff = path_hausdorff_99 = path_hausdorff_95 = path_mean_distance = None
# compute evaluation metrics
evaluate.evaluation(gt_folder,
eval_folder,
evaluation_labels,
mask_dir=mask_folder,
path_dice=os.path.join(eval_folder, 'dice.npy'),
path_hausdorff=path_hausdorff,
path_hausdorff_99=path_hausdorff_99,
path_hausdorff_95=path_hausdorff_95,
path_mean_distance=path_mean_distance,
list_incorrect_labels=list_incorrect_labels,
list_correct_labels=list_correct_labels,
recompute=recompute,
verbose=verbose)
########################################################################################################
# instantiate BrainGenerator object
brain_generator = BrainGenerator(labels_dir=path_label_map,
generation_labels=generation_labels,
output_labels=output_labels,
generation_classes=generation_classes,
prior_distributions=prior_distribution,
prior_means=prior_means,
prior_stds=prior_stds,
output_shape=output_shape)
# create result dir
utils.mkdir(result_dir)
for n in range(n_examples):
# generate new image and corresponding labels
start = time.time()
im, lab = brain_generator.generate_brain()
end = time.time()
print('generation {0:d} took {1:.01f}s'.format(n, end - start))
# save output image and label map
utils.save_volume(np.squeeze(im), brain_generator.aff,
brain_generator.header,
os.path.join(result_dir, 't1_%s.nii.gz' % n))
utils.save_volume(np.squeeze(lab), brain_generator.aff,
brain_generator.header,
os.path.join(result_dir, 't1_labels_%s.nii.gz' % n))
def postprocess_samseg(list_samseg_dir,
list_gt_dir,
path_segmentation_labels,
incorrect_labels,
correct_labels,
list_posteriors_dir=None,
list_thresholds=None,
recompute=False):
""" This function processes the samseg segmentations: it corrects the labels (right/left and 99 to 77), resamples
them to the space of gt_dir, and computes the Dice scores for 1) all_subjects vs. testing subjects only, and 2) all
ROIs vs. lesions only.
It requires that all segmentations are sorted in three subfolders inside samseg_main_dir: t1, flair, and t1_flair.
IMPORTANT: Images are expected to have to following naming convention: <subject_id>.samseg.<contrast>.lesion.mgz,
where <contrast> must either be t1, flair, ***t1_flair***
:param list_samseg_dir: main samseg dir containing the three subfolders t1, flair, t1_flair
:param list_gt_dir: folder with the gt label maps for all subjects
:param path_segmentation_labels: list of segmentation labels
:param incorrect_labels: list of samseg incorrect labels
:param correct_labels: list of labels to correct the wrong one with
:param recompute: whether to recompute files
"""
if list_posteriors_dir is None:
list_posteriors_dir = [None] * len(list_samseg_dir)
for samseg_dir, gt_dir, posteriors_dir, threshold in zip(
list_samseg_dir, list_gt_dir, list_posteriors_dir,
list_thresholds):
# define result directories
samseg_corrected_dir = samseg_dir + '_corrected'
samseg_preprocessed_dir = samseg_dir + '_preprocessed'
if (not os.path.isdir(samseg_preprocessed_dir)) | recompute:
# regroup right/left labels and change 99 to 77
edit_volumes.correct_labels_in_dir(samseg_dir,
incorrect_labels,
correct_labels,
samseg_corrected_dir,
recompute=recompute)
# resample to gt format
edit_volumes.mri_convert_images_in_dir(samseg_corrected_dir,
samseg_preprocessed_dir,
interpolation='nearest',
reference_dir=gt_dir,
recompute=recompute)
# replace lesions by thresholded lesion posteriors
if posteriors_dir is not None:
# resample posteriors to gt format
posteriors_preprocessed_dir = posteriors_dir + '_preprocessed'
edit_volumes.mri_convert_images_in_dir(posteriors_dir,
posteriors_preprocessed_dir,
reference_dir=gt_dir,
recompute=recompute)
# list hard segmentations and posteriors
samseg_postprocessed_dir = samseg_dir + '_postprocessed'
utils.mkdir(samseg_postprocessed_dir)
path_segs = [
path for path in utils.list_images_in_folder(
samseg_preprocessed_dir)
]
path_posteriors = [
path for path in utils.list_images_in_folder(
posteriors_preprocessed_dir)
]
for subject_idx, (path_seg, path_post) in enumerate(
zip(path_segs, path_posteriors)):
path_result = os.path.join(samseg_postprocessed_dir,
os.path.basename(path_seg))
if (not os.path.isfile(path_result)) | recompute:
# replace segmented lesions by thresholded posteriors
seg, aff, h = utils.load_volume(path_seg, im_only=False)
posteriors = utils.load_volume(path_post)
seg[seg == 77] = 2
seg[posteriors > threshold] = 77
utils.save_volume(seg, aff, h, path_result)
else:
samseg_postprocessed_dir = samseg_preprocessed_dir
# compute dice scores with
path_dice_testing = os.path.join(samseg_postprocessed_dir, 'dice.npy')
path_dice_lesions_testing = os.path.join(samseg_postprocessed_dir,
'dice_lesions.npy')
if (not os.path.isfile(path_dice_testing)) | recompute:
dice_evaluation(gt_dir, samseg_postprocessed_dir,
path_segmentation_labels, path_dice_testing)
if (not os.path.isfile(path_dice_lesions_testing)) | recompute:
dice = np.load(path_dice_testing)
np.save(path_dice_lesions_testing, dice[4, :])
downsample = True
# ------------------------------------------------------ Generate ------------------------------------------------------
# instantiate BrainGenerator object
brain_generator = BrainGenerator(labels_dir=path_label_map,
generation_labels=generation_labels,
output_labels=output_labels,
n_neutral_labels=n_neutral_labels,
output_shape=output_shape,
prior_distributions=prior_distributions,
generation_classes=generation_classes,
prior_means=prior_means,
prior_stds=prior_stds,
randomise_res=randomise_res,
data_res=data_res,
thickness=thickness,
downsample=downsample,
blur_range=blur_range)
for n in range(n_examples):
# generate new image and corresponding labels
im, lab = brain_generator.generate_brain()
# save output image and label map
utils.save_volume(im, brain_generator.aff, brain_generator.header,
os.path.join(result_dir, 'image_t1_%s.nii.gz' % n))
utils.save_volume(lab, brain_generator.aff, brain_generator.header,
os.path.join(result_dir, 'labels_t1_%s.nii.gz' % n))
