def run_visualization(input, config, number, output, roi):
"""Utility function to visualize Data Augmentation transformations.
Data augmentation is a key part of the Deep Learning training scheme. This script aims at facilitating the
fine-tuning of data augmentation parameters. To do so, this script provides a step-by-step visualization of the
transformations that are applied on data.
This function applies a series of transformations (defined in a configuration file
``-c``) to ``-n`` 2D slices randomly extracted from an input image (``-i``), and save as png the resulting sample
after each transform.
For example::
ivadomed_visualize_transforms -i t2s.nii.gz -n 1 -c config.json -r t2s_seg.nii.gz
Provides a visualization of a series of three transformation on a randomly selected slice:
.. image:: https://raw.githubusercontent.com/ivadomed/doc-figures/main/scripts/transforms_im.png
:width: 600px
:align: center
And on a binary mask::
ivadomed_visualize_transforms -i t2s_gmseg.nii.gz -n 1 -c config.json -r t2s_seg.nii.gz
Gives:
.. image:: https://raw.githubusercontent.com/ivadomed/doc-figures/main/scripts/transforms_gt.png
:width: 600px
:align: center
Args:
input (string): Image filename. Flag: ``--input``, ``-i``
config (string): Configuration file filename. Flag: ``--config``, ``-c``
number (int): Number of slices randomly extracted. Flag: ``--number``, ``-n``
output (string): Folder path where the results are saved. Flag: ``--ofolder``, ``-o``
roi (string): Filename of the region of interest. Only needed if ROICrop is part of the transformations.
Flag: ``--roi``, ``-r``
"""
# Load context
context = imed_config_manager.ConfigurationManager(config).get_config()
# Create output folder
if not Path(output).is_dir():
Path(output).mkdir(parents=True)
# Slice extracted according to below axis
axis = imed_utils.AXIS_DCT[context[ConfigKW.LOADER_PARAMETERS][LoaderParamsKW.SLICE_AXIS]]
# Get data
input_img, input_data = get_data(input, axis)
# Image or Mask
is_mask = np.array_equal(input_data, input_data.astype(bool))
# Get zooms
zooms = imed_loader_utils.orient_shapes_hwd(input_img.header.get_zooms(), slice_axis=axis)
# Get indexes
indexes = random.sample(range(0, input_data.shape[2]), number)
# Get training transforms
training_transforms, _, _ = imed_transforms.get_subdatasets_transforms(context[ConfigKW.TRANSFORMATION])
if TransformationKW.ROICROP in training_transforms:
if roi and Path(roi).is_file():
roi_img, roi_data = get_data(roi, axis)
else:
raise ValueError("\nPlease provide ROI image (-r) in order to apply ROICrop transformation.")
# Compose transforms
dict_transforms = {}
stg_transforms = ""
for transform_name in training_transforms:
# We skip NumpyToTensor transform since that s only a change of data type
if transform_name == "NumpyToTensor":
continue
# Update stg_transforms
stg_transforms += transform_name + "_"
# Add new transform to Compose
dict_transforms.update({transform_name: training_transforms[transform_name]})
composed_transforms = imed_transforms.Compose(dict_transforms)
# Loop across slices
for i in indexes:
data = [input_data[:, :, i]]
# Init metadata
metadata = SampleMetadata({MetadataKW.ZOOMS: zooms, MetadataKW.DATA_TYPE: "gt" if is_mask else "im"})
# Apply transformations to ROI
if TransformationKW.CENTERCROP in training_transforms or \
(TransformationKW.ROICROP in training_transforms and Path(roi).is_file()):
metadata.__setitem__(MetadataKW.CROP_PARAMS, {})
# Apply transformations to image
stack_im, _ = composed_transforms(sample=data,
metadata=[metadata for _ in range(number)],
data_type="im")
# Plot before / after transformation
fname_out = str(Path(output, stg_transforms + "slice" + str(i) + ".png"))
logger.debug(f"Fname out: {fname_out}.")
logger.debug(f"\t{dict(metadata)}")
# rescale intensities
if len(stg_transforms[:-1].split("_")) == 1:
before = np.rot90(imed_maths.rescale_values_array(data[0], 0.0, 1.0))
else:
before = after
if isinstance(stack_im[0], torch.Tensor):
after = np.rot90(imed_maths.rescale_values_array(stack_im[0].numpy(), 0.0, 1.0))
else:
after = np.rot90(imed_maths.rescale_values_array(stack_im[0], 0.0, 1.0))
# Plot
imed_utils.plot_transformed_sample(before,
after,
list_title=["\n".join(stg_transforms[:-1].split("_")[:-1]),
"\n".join(stg_transforms[:-1].split("_"))],
fname_out=fname_out,
cmap="jet" if is_mask else "gray")
def run_command(context, n_gif=0, thr_increment=None, resume_training=False):
"""Run main command.
This function is central in the ivadomed project as training / testing / evaluation commands are run via this
function. All the process parameters are defined in the config.
Args:
context (dict): Dictionary containing all parameters that are needed for a given process. See
:doc:`configuration_file` for more details.
n_gif (int): Generates a GIF during training if larger than zero, one frame per epoch for a given slice. The
parameter indicates the number of 2D slices used to generate GIFs, one GIF per slice. A GIF shows
predictions of a given slice from the validation sub-dataset. They are saved within the log directory.
thr_increment (float): A threshold analysis is performed at the end of the training using the trained model and
the training + validation sub-dataset to find the optimal binarization threshold. The specified value
indicates the increment between 0 and 1 used during the ROC analysis (e.g. 0.1).
resume_training (bool): Load a saved model ("checkpoint.pth.tar" in the log_directory) for resume training.
This training state is saved everytime a new best model is saved in the log
directory.
Returns:
Float or pandas Dataframe:
If "train" command: Returns floats: best loss score for both training and validation.
If "test" command: Returns a pandas Dataframe: of metrics computed for each subject of the testing
sub dataset and return the prediction metrics before evaluation.
If "segment" command: No return value.
"""
command = copy.deepcopy(context["command"])
log_directory = copy.deepcopy(context["log_directory"])
if not os.path.isdir(log_directory):
print('Creating log directory: {}'.format(log_directory))
os.makedirs(log_directory)
else:
print('Log directory already exists: {}'.format(log_directory))
# Define device
cuda_available, device = imed_utils.define_device(context['gpu'])
# Get subject lists
train_lst, valid_lst, test_lst = imed_loader_utils.get_subdatasets_subjects_list(
context["split_dataset"], context['loader_parameters']['bids_path'],
log_directory)
# Loader params
loader_params = copy.deepcopy(context["loader_parameters"])
if command == "train":
loader_params["contrast_params"]["contrast_lst"] = loader_params[
"contrast_params"]["training_validation"]
else:
loader_params["contrast_params"]["contrast_lst"] = loader_params[
"contrast_params"]["testing"]
if "FiLMedUnet" in context and context["FiLMedUnet"]["applied"]:
loader_params.update(
{"metadata_type": context["FiLMedUnet"]["metadata"]})
# Get transforms for each subdataset
transform_train_params, transform_valid_params, transform_test_params = \
imed_transforms.get_subdatasets_transforms(context["transformation"])
# MODEL PARAMETERS
model_params = copy.deepcopy(context["default_model"])
model_params["folder_name"] = copy.deepcopy(context["model_name"])
model_context_list = [
model_name for model_name in MODEL_LIST
if model_name in context and context[model_name]["applied"]
]
if len(model_context_list) == 1:
model_params["name"] = model_context_list[0]
model_params.update(context[model_context_list[0]])
elif 'Modified3DUNet' in model_context_list and 'FiLMedUnet' in model_context_list and len(
model_context_list) == 2:
model_params["name"] = 'Modified3DUNet'
for i in range(len(model_context_list)):
model_params.update(context[model_context_list[i]])
elif len(model_context_list) > 1:
print(
'ERROR: Several models are selected in the configuration file: {}.'
'Please select only one (i.e. only one where: "applied": true).'.
format(model_context_list))
exit()
model_params['is_2d'] = False if "Modified3DUNet" in model_params[
'name'] else model_params['is_2d']
# Get in_channel from contrast_lst
if loader_params["multichannel"]:
model_params["in_channel"] = len(
loader_params["contrast_params"]["contrast_lst"])
else:
model_params["in_channel"] = 1
# Get out_channel from target_suffix
model_params["out_channel"] = len(loader_params["target_suffix"])
# If multi-class output, then add background class
if model_params["out_channel"] > 1:
model_params.update({"out_channel": model_params["out_channel"] + 1})
# Display for spec' check
imed_utils.display_selected_model_spec(params=model_params)
# Update loader params
if 'object_detection_params' in context:
object_detection_params = context['object_detection_params']
object_detection_params.update({
"gpu":
context['gpu'],
"log_directory":
context['log_directory']
})
loader_params.update(
{"object_detection_params": object_detection_params})
loader_params.update({"model_params": model_params})
# TESTING PARAMS
# Aleatoric uncertainty
if context['uncertainty'][
'aleatoric'] and context['uncertainty']['n_it'] > 0:
transformation_dict = transform_train_params
else:
transformation_dict = transform_test_params
undo_transforms = imed_transforms.UndoCompose(
imed_transforms.Compose(transformation_dict, requires_undo=True))
testing_params = copy.deepcopy(context["training_parameters"])
testing_params.update({'uncertainty': context["uncertainty"]})
testing_params.update({
'target_suffix': loader_params["target_suffix"],
'undo_transforms': undo_transforms,
'slice_axis': loader_params['slice_axis']
})
if command == "train":
imed_utils.display_selected_transfoms(transform_train_params,
dataset_type=["training"])
imed_utils.display_selected_transfoms(transform_valid_params,
dataset_type=["validation"])
elif command == "test":
imed_utils.display_selected_transfoms(transformation_dict,
dataset_type=["testing"])
if command == 'train':
# LOAD DATASET
# Get Validation dataset
ds_valid = imed_loader.load_dataset(**{
**loader_params,
**{
'data_list': valid_lst,
'transforms_params': transform_valid_params,
'dataset_type': 'validation'
}
},
device=device,
cuda_available=cuda_available)
# Get Training dataset
ds_train = imed_loader.load_dataset(**{
**loader_params,
**{
'data_list': train_lst,
'transforms_params': transform_train_params,
'dataset_type': 'training'
}
},
device=device,
cuda_available=cuda_available)
metric_fns = imed_metrics.get_metric_fns(ds_train.task)
# If FiLM, normalize data
if 'film_layers' in model_params and any(model_params['film_layers']):
# Normalize metadata before sending to the FiLM network
results = imed_film.get_film_metadata_models(
ds_train=ds_train,
metadata_type=model_params['metadata'],
debugging=context["debugging"])
ds_train, train_onehotencoder, metadata_clustering_models = results
ds_valid = imed_film.normalize_metadata(
ds_valid, metadata_clustering_models, context["debugging"],
model_params['metadata'])
model_params.update({
"film_onehotencoder":
train_onehotencoder,
"n_metadata":
len([ll for l in train_onehotencoder.categories_ for ll in l])
})
joblib.dump(metadata_clustering_models,
"./" + log_directory + "/clustering_models.joblib")
joblib.dump(train_onehotencoder,
"./" + log_directory + "/one_hot_encoder.joblib")
# Model directory
path_model = os.path.join(log_directory, context["model_name"])
if not os.path.isdir(path_model):
print('Creating model directory: {}'.format(path_model))
os.makedirs(path_model)
if 'film_layers' in model_params and any(
model_params['film_layers']):
joblib.dump(train_onehotencoder,
os.path.join(path_model, "one_hot_encoder.joblib"))
if 'metadata_dict' in ds_train[0]['input_metadata'][0]:
metadata_dict = ds_train[0]['input_metadata'][0][
'metadata_dict']
joblib.dump(
metadata_dict,
os.path.join(path_model, "metadata_dict.joblib"))
else:
print('Model directory already exists: {}'.format(path_model))
# RUN TRAINING
best_training_dice, best_training_loss, best_validation_dice, best_validation_loss = imed_training.train(
model_params=model_params,
dataset_train=ds_train,
dataset_val=ds_valid,
training_params=context["training_parameters"],
log_directory=log_directory,
device=device,
cuda_available=cuda_available,
metric_fns=metric_fns,
n_gif=n_gif,
resume_training=resume_training,
debugging=context["debugging"])
if thr_increment:
# LOAD DATASET
if command != 'train': # If command == train, then ds_valid already load
# Get Validation dataset
ds_valid = imed_loader.load_dataset(**{
**loader_params,
**{
'data_list': valid_lst,
'transforms_params': transform_valid_params,
'dataset_type': 'validation'
}
},
device=device,
cuda_available=cuda_available)
# Get Training dataset with no Data Augmentation
ds_train = imed_loader.load_dataset(**{
**loader_params,
**{
'data_list': train_lst,
'transforms_params': transform_valid_params,
'dataset_type': 'training'
}
},
device=device,
cuda_available=cuda_available)
# Choice of optimisation metric
metric = "recall_specificity" if model_params[
"name"] in imed_utils.CLASSIFIER_LIST else "dice"
# Model path
model_path = os.path.join(log_directory, "best_model.pt")
# Run analysis
thr = imed_testing.threshold_analysis(model_path=model_path,
ds_lst=[ds_train, ds_valid],
model_params=model_params,
testing_params=testing_params,
metric=metric,
increment=thr_increment,
fname_out=os.path.join(
log_directory, "roc.png"),
cuda_available=cuda_available)
# Update threshold in config file
context["postprocessing"]["binarize_prediction"] = {"thr": thr}
if command == 'train':
# Save config file within log_directory and log_directory/model_name
# Done after the threshold_analysis to propate this info in the config files
with open(os.path.join(log_directory, "config_file.json"), 'w') as fp:
json.dump(context, fp, indent=4)
with open(
os.path.join(log_directory, context["model_name"],
context["model_name"] + ".json"), 'w') as fp:
json.dump(context, fp, indent=4)
return best_training_dice, best_training_loss, best_validation_dice, best_validation_loss
if command == 'test':
# LOAD DATASET
ds_test = imed_loader.load_dataset(**{
**loader_params,
**{
'data_list': test_lst,
'transforms_params': transformation_dict,
'dataset_type': 'testing',
'requires_undo': True
}
},
device=device,
cuda_available=cuda_available)
metric_fns = imed_metrics.get_metric_fns(ds_test.task)
if 'film_layers' in model_params and any(model_params['film_layers']):
clustering_path = os.path.join(log_directory,
"clustering_models.joblib")
metadata_clustering_models = joblib.load(clustering_path)
ohe_path = os.path.join(log_directory, "one_hot_encoder.joblib")
one_hot_encoder = joblib.load(ohe_path)
ds_test = imed_film.normalize_metadata(ds_test,
metadata_clustering_models,
context["debugging"],
model_params['metadata'])
model_params.update({
"film_onehotencoder":
one_hot_encoder,
"n_metadata":
len([ll for l in one_hot_encoder.categories_ for ll in l])
})
# RUN INFERENCE
pred_metrics = imed_testing.test(
model_params=model_params,
dataset_test=ds_test,
testing_params=testing_params,
log_directory=log_directory,
device=device,
cuda_available=cuda_available,
metric_fns=metric_fns,
postprocessing=context['postprocessing'])
# RUN EVALUATION
df_results = imed_evaluation.evaluate(
bids_path=loader_params['bids_path'],
log_directory=log_directory,
target_suffix=loader_params["target_suffix"],
eval_params=context["evaluation_parameters"])
return df_results, pred_metrics
if command == 'segment':
bids_ds = bids.BIDS(context["loader_parameters"]["bids_path"])
df = bids_ds.participants.content
subj_lst = df['participant_id'].tolist()
bids_subjects = [
s for s in bids_ds.get_subjects()
if s.record["subject_id"] in subj_lst
]
# Add postprocessing to packaged model
path_model = os.path.join(context['log_directory'],
context['model_name'])
path_model_config = os.path.join(path_model,
context['model_name'] + ".json")
model_config = imed_config_manager.load_json(path_model_config)
model_config['postprocessing'] = context['postprocessing']
with open(path_model_config, 'w') as fp:
json.dump(model_config, fp, indent=4)
options = None
for subject in bids_subjects:
fname_img = subject.record["absolute_path"]
if 'film_layers' in model_params and any(
model_params['film_layers']) and model_params['metadata']:
subj_id = subject.record['subject_id']
metadata = df[df['participant_id'] == subj_id][
model_params['metadata']].values[0]
options = {'metadata': metadata}
pred = imed_inference.segment_volume(path_model,
fname_image=fname_img,
gpu_number=context['gpu'],
options=options)
pred_path = os.path.join(context['log_directory'], "pred_masks")
if not os.path.exists(pred_path):
os.makedirs(pred_path)
filename = subject.record['subject_id'] + "_" + subject.record[
'modality'] + "_pred" + ".nii.gz"
nib.save(pred, os.path.join(pred_path, filename))
def run_command(context, n_gif=0, thr_increment=None, resume_training=False):
"""Run main command.
This function is central in the ivadomed project as training / testing / evaluation commands
are run via this function. All the process parameters are defined in the config.
Args:
context (dict): Dictionary containing all parameters that are needed for a given process. See
:doc:`configuration_file` for more details.
n_gif (int): Generates a GIF during training if larger than zero, one frame per epoch for a given slice. The
parameter indicates the number of 2D slices used to generate GIFs, one GIF per slice. A GIF shows
predictions of a given slice from the validation sub-dataset. They are saved within the output path.
thr_increment (float): A threshold analysis is performed at the end of the training using the trained model and
the training + validation sub-dataset to find the optimal binarization threshold. The specified value
indicates the increment between 0 and 1 used during the ROC analysis (e.g. 0.1).
resume_training (bool): Load a saved model ("checkpoint.pth.tar" in the output directory specified with flag "--path-output" or via the config file "output_path" ' This training state is saved everytime a new best model is saved in the log
argument) for resume training directory.
Returns:
float or pandas.DataFrame or None:
* If "train" command: Returns floats: best loss score for both training and validation.
* If "test" command: Returns a pandas Dataframe: of metrics computed for each subject of
the testing sub-dataset and return the prediction metrics before evaluation.
* If "segment" command: No return value.
"""
command = copy.deepcopy(context["command"])
path_output = set_output_path(context)
# Create a log with the version of the Ivadomed software and the version of the Annexed dataset (if present)
create_dataset_and_ivadomed_version_log(context)
cuda_available, device = imed_utils.define_device(context['gpu_ids'][0])
# BACKWARDS COMPATIBILITY: If bids_path is string, assign to list - Do this here so it propagates to all functions
context['loader_parameters']['path_data'] = imed_utils.format_path_data(
context['loader_parameters']['path_data'])
# Loader params
loader_params = set_loader_params(context, command == "train")
# Get transforms for each subdataset
transform_train_params, transform_valid_params, transform_test_params = \
imed_transforms.get_subdatasets_transforms(context["transformation"])
# MODEL PARAMETERS
model_params, loader_params = set_model_params(context, loader_params)
if command == 'segment':
run_segment_command(context, model_params)
return
# Get subject lists. "segment" command uses all participants of data path, hence no need to split
train_lst, valid_lst, test_lst = imed_loader_utils.get_subdatasets_subjects_list(
context["split_dataset"], context['loader_parameters']['path_data'],
path_output, context["loader_parameters"]['subject_selection'])
# TESTING PARAMS
# Aleatoric uncertainty
if context['uncertainty'][
'aleatoric'] and context['uncertainty']['n_it'] > 0:
transformation_dict = transform_train_params
else:
transformation_dict = transform_test_params
undo_transforms = imed_transforms.UndoCompose(
imed_transforms.Compose(transformation_dict, requires_undo=True))
testing_params = copy.deepcopy(context["training_parameters"])
testing_params.update({'uncertainty': context["uncertainty"]})
testing_params.update({
'target_suffix': loader_params["target_suffix"],
'undo_transforms': undo_transforms,
'slice_axis': loader_params['slice_axis']
})
if command == "train":
imed_utils.display_selected_transfoms(transform_train_params,
dataset_type=["training"])
imed_utils.display_selected_transfoms(transform_valid_params,
dataset_type=["validation"])
elif command == "test":
imed_utils.display_selected_transfoms(transformation_dict,
dataset_type=["testing"])
# Check if multiple raters
check_multiple_raters(command != "train", loader_params)
if command == 'train':
# Get Validation dataset
ds_valid = get_dataset(loader_params, valid_lst,
transform_valid_params, cuda_available, device,
'validation')
# Get Training dataset
ds_train = get_dataset(loader_params, train_lst,
transform_train_params, cuda_available, device,
'training')
metric_fns = imed_metrics.get_metric_fns(ds_train.task)
# If FiLM, normalize data
if 'film_layers' in model_params and any(model_params['film_layers']):
model_params, ds_train, ds_valid, train_onehotencoder = \
film_normalize_data(context, model_params, ds_train, ds_valid, path_output)
else:
train_onehotencoder = None
# Model directory
create_path_model(context, model_params, ds_train, path_output,
train_onehotencoder)
save_config_file(context, path_output)
# RUN TRAINING
best_training_dice, best_training_loss, best_validation_dice, best_validation_loss = imed_training.train(
model_params=model_params,
dataset_train=ds_train,
dataset_val=ds_valid,
training_params=context["training_parameters"],
path_output=path_output,
device=device,
cuda_available=cuda_available,
metric_fns=metric_fns,
n_gif=n_gif,
resume_training=resume_training,
debugging=context["debugging"])
if thr_increment:
# LOAD DATASET
if command != 'train': # If command == train, then ds_valid already load
# Get Validation dataset
ds_valid = get_dataset(loader_params, valid_lst,
transform_valid_params, cuda_available,
device, 'validation')
# Get Training dataset with no Data Augmentation
ds_train = get_dataset(loader_params, train_lst,
transform_valid_params, cuda_available, device,
'training')
# Choice of optimisation metric
metric = "recall_specificity" if model_params[
"name"] in imed_utils.CLASSIFIER_LIST else "dice"
# Model path
model_path = os.path.join(path_output, "best_model.pt")
# Run analysis
thr = imed_testing.threshold_analysis(model_path=model_path,
ds_lst=[ds_train, ds_valid],
model_params=model_params,
testing_params=testing_params,
metric=metric,
increment=thr_increment,
fname_out=os.path.join(
path_output, "roc.png"),
cuda_available=cuda_available)
# Update threshold in config file
context["postprocessing"]["binarize_prediction"] = {"thr": thr}
save_config_file(context, path_output)
if command == 'train':
return best_training_dice, best_training_loss, best_validation_dice, best_validation_loss
if command == 'test':
# LOAD DATASET
ds_test = imed_loader.load_dataset(**{
**loader_params,
**{
'data_list': test_lst,
'transforms_params': transformation_dict,
'dataset_type': 'testing',
'requires_undo': True
}
},
device=device,
cuda_available=cuda_available)
metric_fns = imed_metrics.get_metric_fns(ds_test.task)
if 'film_layers' in model_params and any(model_params['film_layers']):
ds_test, model_params = update_film_model_params(
context, ds_test, model_params, path_output)
# RUN INFERENCE
pred_metrics = imed_testing.test(
model_params=model_params,
dataset_test=ds_test,
testing_params=testing_params,
path_output=path_output,
device=device,
cuda_available=cuda_available,
metric_fns=metric_fns,
postprocessing=context['postprocessing'])
# RUN EVALUATION
df_results = imed_evaluation.evaluate(
path_data=loader_params['path_data'],
path_output=path_output,
target_suffix=loader_params["target_suffix"],
eval_params=context["evaluation_parameters"])
return df_results, pred_metrics
def segment_volume(folder_model, fname_images, gpu_number=0, options=None):
"""Segment an image.
Segment an image (`fname_image`) using a pre-trained model (`folder_model`). If provided, a region of interest
(`fname_roi`) is used to crop the image prior to segment it.
Args:
folder_model (str): foldername which contains
(1) the model ('folder_model/folder_model.pt') to use
(2) its configuration file ('folder_model/folder_model.json') used for the training,
see https://github.com/neuropoly/ivadomed/wiki/configuration-file
fname_images (list): list of image filenames (e.g. .nii.gz) to segment. Multichannel models require multiple
images to segment, e.i., len(fname_images) > 1.
gpu_number (int): Number representing gpu number if available.
options (dict): Contains postprocessing steps and prior filename (fname_prior) which is an image filename
(e.g., .nii.gz) containing processing information (e.i., spinal cord segmentation, spinal location or MS
lesion classification)
e.g., spinal cord centerline, used to crop the image prior to segment it if provided.
The segmentation is not performed on the slices that are empty in this image.
Returns:
list: List of nibabel objects containing the soft segmentation(s), one per prediction class.
list: List of target suffix associated with each prediction in `pred_list`
"""
# Define device
cuda_available = torch.cuda.is_available()
device = torch.device("cpu") if not cuda_available else torch.device(
"cuda:" + str(gpu_number))
# Check if model folder exists and get filenames
fname_model, fname_model_metadata = imed_models.get_model_filenames(
folder_model)
# Load model training config
context = imed_config_manager.ConfigurationManager(
fname_model_metadata).get_config()
postpro_list = [
'binarize_prediction', 'keep_largest', ' fill_holes', 'remove_small'
]
if options is not None and any(pp in options for pp in postpro_list):
postpro = {}
if 'binarize_prediction' in options and options['binarize_prediction']:
postpro['binarize_prediction'] = {
"thr": options['binarize_prediction']
}
if 'keep_largest' in options and options['keep_largest'] is not None:
if options['keep_largest']:
postpro['keep_largest'] = {}
# Remove key in context if value set to 0
elif 'keep_largest' in context['postprocessing']:
del context['postprocessing']['keep_largest']
if 'fill_holes' in options and options['fill_holes'] is not None:
if options['fill_holes']:
postpro['fill_holes'] = {}
# Remove key in context if value set to 0
elif 'fill_holes' in context['postprocessing']:
del context['postprocessing']['fill_holes']
if 'remove_small' in options and options['remove_small'] and \
('mm' in options['remove_small'][-1] or 'vox' in options['remove_small'][-1]):
unit = 'mm3' if 'mm3' in options['remove_small'][-1] else 'vox'
thr = [int(t.replace(unit, "")) for t in options['remove_small']]
postpro['remove_small'] = {"unit": unit, "thr": thr}
context['postprocessing'].update(postpro)
# LOADER
loader_params = context["loader_parameters"]
slice_axis = imed_utils.AXIS_DCT[loader_params['slice_axis']]
metadata = {}
fname_roi = None
fname_prior = options['fname_prior'] if (options is not None) and (
'fname_prior' in options) else None
if fname_prior is not None:
if 'roi_params' in loader_params and loader_params['roi_params'][
'suffix'] is not None:
fname_roi = fname_prior
# TRANSFORMATIONS
# If ROI is not provided then force center cropping
if fname_roi is None and 'ROICrop' in context["transformation"].keys():
print(
"\n WARNING: fname_roi has not been specified, then a cropping around the center of the image is "
"performed instead of a cropping around a Region of Interest.")
context["transformation"] = dict(
(key, value) if key != 'ROICrop' else ('CenterCrop', value)
for (key, value) in context["transformation"].items())
if 'object_detection_params' in context and \
context['object_detection_params']['object_detection_path'] is not None:
imed_obj_detect.bounding_box_prior(
fname_prior, metadata, slice_axis,
context['object_detection_params']['safety_factor'])
metadata = [metadata] * len(fname_images)
# Compose transforms
_, _, transform_test_params = imed_transforms.get_subdatasets_transforms(
context["transformation"])
tranform_lst, undo_transforms = imed_transforms.prepare_transforms(
transform_test_params)
# Force filter_empty_mask to False if fname_roi = None
if fname_roi is None and 'filter_empty_mask' in loader_params[
"slice_filter_params"]:
print(
"\nWARNING: fname_roi has not been specified, then the entire volume is processed."
)
loader_params["slice_filter_params"]["filter_empty_mask"] = False
filename_pairs = [(fname_images, None, fname_roi,
metadata if isinstance(metadata, list) else [metadata])]
kernel_3D = bool('Modified3DUNet' in context and context['Modified3DUNet']['applied']) or \
not context['default_model']['is_2d']
if kernel_3D:
ds = imed_loader.MRI3DSubVolumeSegmentationDataset(
filename_pairs,
transform=tranform_lst,
length=context["Modified3DUNet"]["length_3D"],
stride=context["Modified3DUNet"]["stride_3D"])
else:
ds = imed_loader.MRI2DSegmentationDataset(
filename_pairs,
slice_axis=slice_axis,
cache=True,
transform=tranform_lst,
slice_filter_fn=imed_loader_utils.SliceFilter(
**loader_params["slice_filter_params"]))
ds.load_filenames()
if kernel_3D:
print("\nLoaded {} {} volumes of shape {}.".format(
len(ds), loader_params['slice_axis'],
context['Modified3DUNet']['length_3D']))
else:
print("\nLoaded {} {} slices.".format(len(ds),
loader_params['slice_axis']))
model_params = {}
if 'FiLMedUnet' in context and context['FiLMedUnet']['applied']:
metadata_dict = joblib.load(
os.path.join(folder_model, 'metadata_dict.joblib'))
for idx in ds.indexes:
for i in range(len(idx)):
idx[i]['input_metadata'][0][context['FiLMedUnet']
['metadata']] = options['metadata']
idx[i]['input_metadata'][0]['metadata_dict'] = metadata_dict
ds = imed_film.normalize_metadata(ds, None, context["debugging"],
context['FiLMedUnet']['metadata'])
onehotencoder = joblib.load(
os.path.join(folder_model, 'one_hot_encoder.joblib'))
model_params.update({
"name":
'FiLMedUnet',
"film_onehotencoder":
onehotencoder,
"n_metadata":
len([ll for l in onehotencoder.categories_ for ll in l])
})
# Data Loader
data_loader = DataLoader(
ds,
batch_size=context["training_parameters"]["batch_size"],
shuffle=False,
pin_memory=True,
collate_fn=imed_loader_utils.imed_collate,
num_workers=0)
# MODEL
if fname_model.endswith('.pt'):
model = torch.load(fname_model, map_location=device)
# Inference time
model.eval()
# Loop across batches
preds_list, slice_idx_list = [], []
last_sample_bool, volume, weight_matrix = False, None, None
for i_batch, batch in enumerate(data_loader):
with torch.no_grad():
img = imed_utils.cuda(batch['input'],
cuda_available=cuda_available)
if ('FiLMedUnet' in context and context['FiLMedUnet']['applied']) or \
('HeMISUnet' in context and context['HeMISUnet']['applied']):
metadata = imed_training.get_metadata(batch["input_metadata"],
model_params)
preds = model(img, metadata)
else:
preds = model(img) if fname_model.endswith(
'.pt') else onnx_inference(fname_model, img)
preds = preds.cpu()
# Set datatype to gt since prediction should be processed the same way as gt
for b in batch['input_metadata']:
for modality in b:
modality['data_type'] = 'gt'
# Reconstruct 3D object
for i_slice in range(len(preds)):
if "bounding_box" in batch['input_metadata'][i_slice][0]:
imed_obj_detect.adjust_undo_transforms(
undo_transforms.transforms, batch, i_slice)
batch['gt_metadata'] = [[metadata[0]] * preds.shape[1]
for metadata in batch['input_metadata']]
if kernel_3D:
preds_undo, metadata, last_sample_bool, volume, weight_matrix = \
volume_reconstruction(batch, preds, undo_transforms, i_slice, volume, weight_matrix)
preds_list = [np.array(preds_undo)]
else:
# undo transformations
preds_i_undo, metadata_idx = undo_transforms(
preds[i_slice],
batch["input_metadata"][i_slice],
data_type='gt')
# Add new segmented slice to preds_list
preds_list.append(np.array(preds_i_undo))
# Store the slice index of preds_i_undo in the original 3D image
slice_idx_list.append(
int(batch['input_metadata'][i_slice][0]['slice_index']))
# If last batch and last sample of this batch, then reconstruct 3D object
if (i_batch == len(data_loader) - 1
and i_slice == len(batch['gt']) - 1) or last_sample_bool:
pred_nib = pred_to_nib(
data_lst=preds_list,
fname_ref=fname_images[0],
fname_out=None,
z_lst=slice_idx_list,
slice_axis=slice_axis,
kernel_dim='3d' if kernel_3D else '2d',
debug=False,
bin_thr=-1,
postprocessing=context['postprocessing'])
pred_list = split_classes(pred_nib)
target_list = context['loader_parameters']['target_suffix']
return pred_list, target_list
def run_visualization(input, config, number, output, roi):
"""Utility function to visualize Data Augmentation transformations.
Data augmentation is a key part of the Deep Learning training scheme. This script aims at facilitating the
fine-tuning of data augmentation parameters. To do so, this script provides a step-by-step visualization of the
transformations that are applied on data.
This function applies a series of transformations (defined in a configuration file
``-c``) to ``-n`` 2D slices randomly extracted from an input image (``-i``), and save as png the resulting sample
after each transform.
For example::
ivadomed_visualize_transforms -i t2s.nii.gz -n 1 -c config.json -r t2s_seg.nii.gz
Provides a visualization of a series of three transformation on a randomly selected slice:
.. image:: ../../images/transforms_im.png
:width: 600px
:align: center
And on a binary mask::
ivadomed_visualize_transforms -i t2s_gmseg.nii.gz -n 1 -c config.json -r t2s_seg.nii.gz
Gives:
.. image:: ../../images/transforms_gt.png
:width: 600px
:align: center
Args:
input (string): Image filename. Flag: ``--input``, ``-i``
config (string): Configuration file filename. Flag: ``--config``, ``-c``
number (int): Number of slices randomly extracted. Flag: ``--number``, ``-n``
output (string): Folder path where the results are saved. Flag: ``--ofolder``, ``-o``
roi (string): Filename of the region of interest. Only needed if ROICrop is part of the transformations.
Flag: ``--roi``, ``-r``
"""
# Load context
with open(config, "r") as fhandle:
context = json.load(fhandle)
# Create output folder
if not os.path.isdir(output):
os.makedirs(output)
# Slice extracted according to below axis
axis = imed_utils.AXIS_DCT[context["loader_parameters"]["slice_axis"]]
# Get data
input_img, input_data = get_data(input, axis)
# Image or Mask
is_mask = np.array_equal(input_data, input_data.astype(bool))
# Get zooms
zooms = imed_loader_utils.orient_shapes_hwd(input_img.header.get_zooms(),
slice_axis=axis)
# Get indexes
indexes = random.sample(range(0, input_data.shape[2]), number)
# Get training transforms
training_transforms, _, _ = imed_transforms.get_subdatasets_transforms(
context["transformation"])
if "ROICrop" in training_transforms:
if roi and os.path.isfile(roi):
roi_img, roi_data = get_data(roi, axis)
else:
print(
"\nPlease provide ROI image (-r) in order to apply ROICrop transformation."
)
exit()
# Compose transforms
dict_transforms = {}
stg_transforms = ""
for transform_name in training_transforms:
# We skip NumpyToTensor transform since that s only a change of data type
if transform_name == "NumpyToTensor":
continue
# Update stg_transforms
stg_transforms += transform_name + "_"
# Add new transform to Compose
dict_transforms.update(
{transform_name: training_transforms[transform_name]})
composed_transforms = imed_transforms.Compose(dict_transforms)
# Loop across slices
for i in indexes:
data = [input_data[:, :, i]]
# Init metadata
metadata = imed_loader_utils.SampleMetadata({
"zooms":
zooms,
"data_type":
"gt" if is_mask else "im"
})
# Apply transformations to ROI
if "CenterCrop" in training_transforms or (
"ROICrop" in training_transforms and os.path.isfile(roi)):
metadata.__setitem__('crop_params', {})
# Apply transformations to image
stack_im, _ = composed_transforms(
sample=data,
metadata=[metadata for _ in range(number)],
data_type="im")
# Plot before / after transformation
fname_out = os.path.join(
output, stg_transforms + "slice" + str(i) + ".png")
print("Fname out: {}.".format(fname_out))
print("\t{}".format(dict(metadata)))
# rescale intensities
if len(stg_transforms[:-1].split("_")) == 1:
before = np.rot90(
imed_maths.rescale_values_array(data[0], 0.0, 1.0))
else:
before = after
after = np.rot90(
imed_maths.rescale_values_array(stack_im[0], 0.0, 1.0))
# Plot
imed_utils.plot_transformed_sample(
before,
after,
list_title=[
"\n".join(stg_transforms[:-1].split("_")[:-1]),
"\n".join(stg_transforms[:-1].split("_"))
],
fname_out=fname_out,
cmap="jet" if is_mask else "gray")
def segment_volume(folder_model: str,
fname_images: list,
gpu_id: int = 0,
options: dict = None):
"""Segment an image.
Segment an image (`fname_image`) using a pre-trained model (`folder_model`). If provided, a region of interest
(`fname_roi`) is used to crop the image prior to segment it.
Args:
folder_model (str): foldername which contains
(1) the model ('folder_model/folder_model.pt') to use
(2) its configuration file ('folder_model/folder_model.json') used for the training,
see https://github.com/neuropoly/ivadomed/wiki/configuration-file
fname_images (list): list of image filenames (e.g. .nii.gz) to segment. Multichannel models require multiple
images to segment, e.i., len(fname_images) > 1.
gpu_id (int): Number representing gpu number if available. Currently does NOT support multiple GPU segmentation.
options (dict): Contains postprocessing steps and prior filename (fname_prior) which is an image filename
(e.g., .nii.gz) containing processing information (e.i., spinal cord segmentation, spinal location or MS
lesion classification)
e.g., spinal cord centerline, used to crop the image prior to segment it if provided.
The segmentation is not performed on the slices that are empty in this image.
Returns:
list: List of nibabel objects containing the soft segmentation(s), one per prediction class.
list: List of target suffix associated with each prediction in `pred_list`
"""
# Check if model folder exists and get filenames to be stored as string
fname_model: str
fname_model_metadata: str
fname_model, fname_model_metadata = imed_models.get_model_filenames(
folder_model)
# Load model training config
context = imed_config_manager.ConfigurationManager(
fname_model_metadata).get_config()
postpro_list = [
'binarize_prediction', 'keep_largest', ' fill_holes', 'remove_small'
]
if options is not None and any(pp in options for pp in postpro_list):
set_postprocessing_options(options, context)
# LOADER
loader_params = context["loader_parameters"]
slice_axis = imed_utils.AXIS_DCT[loader_params['slice_axis']]
metadata = {}
fname_roi = None
fname_prior = options['fname_prior'] if (options is not None) and (
'fname_prior' in options) else None
if fname_prior is not None:
if 'roi_params' in loader_params and loader_params['roi_params'][
'suffix'] is not None:
fname_roi = fname_prior
# TRANSFORMATIONS
metadata = process_transformations(context, fname_roi, fname_prior,
metadata, slice_axis, fname_images)
# Compose transforms
_, _, transform_test_params = imed_transforms.get_subdatasets_transforms(
context["transformation"])
tranform_lst, undo_transforms = imed_transforms.prepare_transforms(
transform_test_params)
# Force filter_empty_mask to False if fname_roi = None
if fname_roi is None and 'filter_empty_mask' in loader_params[
"slice_filter_params"]:
logger.warning(
"fname_roi has not been specified, then the entire volume is processed."
)
loader_params["slice_filter_params"]["filter_empty_mask"] = False
filename_pairs = [(fname_images, None, fname_roi,
metadata if isinstance(metadata, list) else [metadata])]
kernel_3D = bool('Modified3DUNet' in context and context['Modified3DUNet']['applied']) or \
not context['default_model']['is_2d']
if kernel_3D:
ds = imed_loader.MRI3DSubVolumeSegmentationDataset(
filename_pairs,
transform=tranform_lst,
length=context["Modified3DUNet"]["length_3D"],
stride=context["Modified3DUNet"]["stride_3D"])
logger.info(
f"Loaded {len(ds)} {loader_params['slice_axis']} volumes of shape "
f"{context['Modified3DUNet']['length_3D']}.")
else:
ds = imed_loader.MRI2DSegmentationDataset(
filename_pairs,
slice_axis=slice_axis,
cache=True,
transform=tranform_lst,
slice_filter_fn=imed_loader_utils.SliceFilter(
**loader_params["slice_filter_params"]))
ds.load_filenames()
logger.info(f"Loaded {len(ds)} {loader_params['slice_axis']} slices.")
model_params = {}
if 'FiLMedUnet' in context and context['FiLMedUnet']['applied']:
onehotencoder = get_onehotencoder(context, folder_model, options, ds)
model_params.update({
"name":
'FiLMedUnet',
"film_onehotencoder":
onehotencoder,
"n_metadata":
len([ll for l in onehotencoder.categories_ for ll in l])
})
# Data Loader
data_loader = DataLoader(
ds,
batch_size=context["training_parameters"]["batch_size"],
shuffle=False,
pin_memory=True,
collate_fn=imed_loader_utils.imed_collate,
num_workers=0)
# Loop across batches
preds_list, slice_idx_list = [], []
last_sample_bool, weight_matrix, volume = False, None, None
for i_batch, batch in enumerate(data_loader):
preds = get_preds(context, fname_model, model_params, gpu_id, batch)
# Set datatype to gt since prediction should be processed the same way as gt
for b in batch['input_metadata']:
for modality in b:
modality['data_type'] = 'gt'
# Reconstruct 3D object
pred_list, target_list, last_sample_bool, weight_matrix, volume = reconstruct_3d_object(
context, batch, undo_transforms, preds, preds_list, kernel_3D,
slice_axis, slice_idx_list, data_loader, fname_images, i_batch,
last_sample_bool, weight_matrix, volume)
return pred_list, target_list
"contrast_params"]["training_validation"]
else:
loader_params["contrast_params"]["contrast_lst"] = loader_params[
"contrast_params"]["testing"]
if "FiLMedUnet" in context and context["FiLMedUnet"]["applied"]:
loader_params.update({"metadata_type": context["FiLMedUnet"]["metadata"]})
# Load metadata necessary to balance the loader
if context['training_parameters']['balance_samples']['applied'] and \
context['training_parameters']['balance_samples']['type'] != 'gt':
loader_params.update({
"metadata_type":
context['training_parameters']['balance_samples']['type']
})
# Get transforms for each subdataset
transform_train_params, transform_valid_params, transform_test_params = \
imed_transforms.get_subdatasets_transforms(context["transformation"])
# MODEL PARAMETERS
model_params = copy.deepcopy(context["default_model"])
model_params["folder_name"] = copy.deepcopy(context["model_name"])
model_context_list = [
model_name for model_name in MODEL_LIST
if model_name in context and context[model_name]["applied"]
]
if len(model_context_list) == 1:
model_params["name"] = model_context_list[0]
model_params.update(context[model_context_list[0]])
elif 'Modified3DUNet' in model_context_list and 'FiLMedUnet' in model_context_list and len(
model_context_list) == 2:
model_params["name"] = 'Modified3DUNet'
for i in range(len(model_context_list)):
model_params.update(context[model_context_list[i]])
def segment_volume(folder_model: str,
fname_images: list,
gpu_id: int = 0,
options: dict = None):
"""Segment an image.
Segment an image (`fname_image`) using a pre-trained model (`folder_model`). If provided, a region of interest
(`fname_roi`) is used to crop the image prior to segment it.
Args:
folder_model (str): foldername which contains
(1) the model ('folder_model/folder_model.pt') to use
(2) its configuration file ('folder_model/folder_model.json') used for the training,
see https://github.com/neuropoly/ivadomed/wiki/configuration-file
fname_images (list): list of image filenames (e.g. .nii.gz) to segment. Multichannel models require multiple
images to segment, e.i., len(fname_images) > 1.
gpu_id (int): Number representing gpu number if available. Currently does NOT support multiple GPU segmentation.
options (dict): This can optionally contain any of the following key-value pairs:
* 'binarize_prediction': (float) Binarize segmentation with specified threshold. \
Predictions below the threshold become 0, and predictions above or equal to \
threshold become 1. Set to -1 for no thresholding (i.e., soft segmentation).
* 'binarize_maxpooling': (bool) Binarize by setting to 1 the voxel having the maximum prediction across \
all classes. Useful for multiclass models.
* 'fill_holes': (bool) Fill small holes in the segmentation.
* 'keep_largest': (bool) Keep the largest connected-object for each class from the output segmentation.
* 'remove_small': (list of str) Minimal object size to keep with unit (mm3 or vox). A single value can be provided \
or one value per prediction class. Single value example: ["1mm3"], ["5vox"]. Multiple values \
example: ["10", "20", "10vox"] (remove objects smaller than 10 voxels for class 1 and 3, \
and smaller than 20 voxels for class 2).
* 'pixel_size': (list of float) List of microscopy pixel size in micrometers. \
Length equals 2 [PixelSizeX, PixelSizeY] for 2D or 3 [PixelSizeX, PixelSizeY, PixelSizeZ] for 3D, \
where X is the width, Y the height and Z the depth of the image.
* 'pixel_size_units': (str) Units of pixel size (Must be either "mm", "um" or "nm")
* 'overlap_2D': (list of int) List of overlaps in pixels for 2D patching. Length equals 2 [OverlapX, OverlapY], \
where X is the width and Y the height of the image.
* 'metadata': (str) Film metadata.
* 'fname_prior': (str) An image filename (e.g., .nii.gz) containing processing information \
(e.g., spinal cord segmentation, spinal location or MS lesion classification, spinal cord centerline), \
used to crop the image prior to segment it if provided. \
The segmentation is not performed on the slices that are empty in this image.
Returns:
list, list: List of nibabel objects containing the soft segmentation(s), one per prediction class, \
List of target suffix associated with each prediction in `pred_list`
"""
# Check if model folder exists and get filenames to be stored as string
fname_model: str
fname_model_metadata: str
fname_model, fname_model_metadata = imed_models.get_model_filenames(
folder_model)
# Load model training config
context = imed_config_manager.ConfigurationManager(
fname_model_metadata).get_config()
postpro_list = [
'binarize_prediction', 'binarize_maxpooling', 'keep_largest',
' fill_holes', 'remove_small'
]
if options is not None and any(pp in options for pp in postpro_list):
set_postprocessing_options(options, context)
# LOADER
loader_params = context[ConfigKW.LOADER_PARAMETERS]
slice_axis = imed_utils.AXIS_DCT[loader_params[LoaderParamsKW.SLICE_AXIS]]
metadata = {}
fname_roi = None
if (options is not None) and (OptionKW.FNAME_PRIOR in options):
fname_prior = options.get(OptionKW.FNAME_PRIOR)
else:
fname_prior = None
if fname_prior is not None:
if LoaderParamsKW.ROI_PARAMS in loader_params and loader_params[
LoaderParamsKW.ROI_PARAMS][ROIParamsKW.SUFFIX] is not None:
fname_roi = fname_prior
# TRANSFORMATIONS
metadata = process_transformations(context, fname_roi, fname_prior,
metadata, slice_axis, fname_images)
# Compose transforms
_, _, transform_test_params = imed_transforms.get_subdatasets_transforms(
context[ConfigKW.TRANSFORMATION])
tranform_lst, undo_transforms = imed_transforms.prepare_transforms(
transform_test_params)
# Force filter_empty_mask to False if fname_roi = None
if fname_roi is None and SliceFilterParamsKW.FILTER_EMPTY_MASK in loader_params[
LoaderParamsKW.SLICE_FILTER_PARAMS]:
logger.warning(
"fname_roi has not been specified, then the entire volume is processed."
)
loader_params[LoaderParamsKW.SLICE_FILTER_PARAMS][
SliceFilterParamsKW.FILTER_EMPTY_MASK] = False
kernel_3D = bool(ConfigKW.MODIFIED_3D_UNET in context and context[ConfigKW.MODIFIED_3D_UNET][ModelParamsKW.APPLIED]) or \
not context[ConfigKW.DEFAULT_MODEL][ModelParamsKW.IS_2D]
# Assign length_2D and stride_2D for 2D patching
length_2D = context[ConfigKW.DEFAULT_MODEL][ModelParamsKW.LENGTH_2D] if \
ModelParamsKW.LENGTH_2D in context[ConfigKW.DEFAULT_MODEL] else []
stride_2D = context[ConfigKW.DEFAULT_MODEL][ModelParamsKW.STRIDE_2D] if \
ModelParamsKW.STRIDE_2D in context[ConfigKW.DEFAULT_MODEL] else []
is_2d_patch = bool(length_2D)
if is_2d_patch and (options is not None) and (OptionKW.OVERLAP_2D
in options):
overlap_2D = options.get(OptionKW.OVERLAP_2D)
# Swap OverlapX and OverlapY resulting in an array in order [OverlapY, OverlapX]
# to match length_2D and stride_2D in [Height, Width] orientation.
overlap_2D[1], overlap_2D[0] = overlap_2D[0], overlap_2D[1]
# Adjust stride_2D with overlap_2D
stride_2D = [x1 - x2 for (x1, x2) in zip(length_2D, overlap_2D)]
# Add microscopy pixel size and pixel size units from options to metadata for filenames_pairs
if (options is not None) and (OptionKW.PIXEL_SIZE in options):
metadata[MetadataKW.PIXEL_SIZE] = options.get(OptionKW.PIXEL_SIZE)
if (options is not None) and (OptionKW.PIXEL_SIZE_UNITS in options):
metadata[MetadataKW.PIXEL_SIZE_UNITS] = options.get(
OptionKW.PIXEL_SIZE_UNITS)
filename_pairs = [(fname_images, None, fname_roi,
metadata if isinstance(metadata, list) else [metadata])]
if kernel_3D:
ds = MRI3DSubVolumeSegmentationDataset(
filename_pairs,
transform=tranform_lst,
length=context[ConfigKW.MODIFIED_3D_UNET][ModelParamsKW.LENGTH_3D],
stride=context[ConfigKW.MODIFIED_3D_UNET][ModelParamsKW.STRIDE_3D],
slice_axis=slice_axis)
logger.info(
f"Loaded {len(ds)} {loader_params[LoaderParamsKW.SLICE_AXIS]} volumes of shape "
f"{context[ConfigKW.MODIFIED_3D_UNET][ModelParamsKW.LENGTH_3D]}.")
else:
ds = MRI2DSegmentationDataset(
filename_pairs,
length=length_2D,
stride=stride_2D,
slice_axis=slice_axis,
nibabel_cache=True,
transform=tranform_lst,
slice_filter_fn=SliceFilter(
**loader_params[LoaderParamsKW.SLICE_FILTER_PARAMS]))
ds.load_filenames()
if is_2d_patch:
logger.info(
f"Loaded {len(ds)} {loader_params[LoaderParamsKW.SLICE_AXIS]} patches of shape {length_2D}."
)
else:
logger.info(
f"Loaded {len(ds)} {loader_params[LoaderParamsKW.SLICE_AXIS]} slices."
)
model_params = {}
if ConfigKW.FILMED_UNET in context and context[ConfigKW.FILMED_UNET][
ModelParamsKW.APPLIED]:
onehotencoder = get_onehotencoder(context, folder_model, options, ds)
model_params.update({
ModelParamsKW.NAME:
ConfigKW.FILMED_UNET,
ModelParamsKW.FILM_ONEHOTENCODER:
onehotencoder,
ModelParamsKW.N_METADATA:
len([ll for l in onehotencoder.categories_ for ll in l])
})
# Data Loader
data_loader = DataLoader(ds,
batch_size=context[ConfigKW.TRAINING_PARAMETERS][
TrainingParamsKW.BATCH_SIZE],
shuffle=False,
pin_memory=True,
collate_fn=imed_loader_utils.imed_collate,
num_workers=0)
# Loop across batches
preds_list, slice_idx_list = [], []
last_sample_bool, weight_matrix, volume, image = False, None, None, None
for i_batch, batch in enumerate(data_loader):
preds = get_preds(context, fname_model, model_params, gpu_id, batch)
# Set datatype to gt since prediction should be processed the same way as gt
for b in batch[MetadataKW.INPUT_METADATA]:
for modality in b:
modality['data_type'] = 'gt'
# Reconstruct 3D object
pred_list, target_list, last_sample_bool, weight_matrix, volume, image = reconstruct_3d_object(
context, batch, undo_transforms, preds, preds_list, kernel_3D,
is_2d_patch, slice_axis, slice_idx_list, data_loader, fname_images,
i_batch, last_sample_bool, weight_matrix, volume, image)
return pred_list, target_list
def segment_volume(folder_model, fname_image, fname_prior=None, gpu_number=0):
"""Segment an image.
Segment an image (`fname_image`) using a pre-trained model (`folder_model`). If provided, a region of interest
(`fname_roi`) is used to crop the image prior to segment it.
Args:
folder_model (str): foldername which contains
(1) the model ('folder_model/folder_model.pt') to use
(2) its configuration file ('folder_model/folder_model.json') used for the training,
see https://github.com/neuropoly/ivadomed/wiki/configuration-file
fname_image (str): image filename (e.g. .nii.gz) to segment.
fname_prior (str): Image filename (e.g. .nii.gz) containing processing information (e.i. spinal cord
segmentation, spinal location or MS lesion classification)
e.g. spinal cord centerline, used to crop the image prior to segment it if provided.
The segmentation is not performed on the slices that are empty in this image.
gpu_number (int): Number representing gpu number if available.
Returns:
nibabelObject: Object containing the soft segmentation.
"""
# Define device
cuda_available = torch.cuda.is_available()
device = torch.device("cpu") if not cuda_available else torch.device(
"cuda:" + str(gpu_number))
# Check if model folder exists and get filenames
fname_model, fname_model_metadata = imed_models.get_model_filenames(
folder_model)
# Load model training config
with open(fname_model_metadata, "r") as fhandle:
context = json.load(fhandle)
# LOADER
loader_params = context["loader_parameters"]
slice_axis = AXIS_DCT[loader_params['slice_axis']]
metadata = {}
fname_roi = None
if fname_prior is not None:
if 'roi_params' in loader_params and loader_params['roi_params'][
'suffix'] is not None:
fname_roi = fname_prior
# TRANSFORMATIONS
# If ROI is not provided then force center cropping
if fname_roi is None and 'ROICrop' in context["transformation"].keys():
print(
"\nWARNING: fname_roi has not been specified, then a cropping around the center of the image is performed"
" instead of a cropping around a Region of Interest.")
context["transformation"] = dict(
(key, value) if key != 'ROICrop' else ('CenterCrop', value)
for (key, value) in context["transformation"].items())
if 'object_detection_params' in context and \
context['object_detection_params']['object_detection_path'] is not None:
imed_obj_detect.bounding_box_prior(fname_prior, metadata,
slice_axis)
# Compose transforms
_, _, transform_test_params = imed_transforms.get_subdatasets_transforms(
context["transformation"])
tranform_lst, undo_transforms = imed_transforms.prepare_transforms(
transform_test_params)
# Force filter_empty_mask to False if fname_roi = None
if fname_roi is None and 'filter_empty_mask' in loader_params[
"slice_filter_params"]:
print(
"\nWARNING: fname_roi has not been specified, then the entire volume is processed."
)
loader_params["slice_filter_params"]["filter_empty_mask"] = False
filename_pairs = [([fname_image], None, fname_roi, [metadata])]
kernel_3D = bool('UNet3D' in context and context['UNet3D']['applied'])
if kernel_3D:
ds = imed_loader.MRI3DSubVolumeSegmentationDataset(
filename_pairs,
transform=tranform_lst,
length=context["UNet3D"]["length_3D"],
stride=context["UNet3D"]["stride_3D"])
else:
ds = imed_loader.MRI2DSegmentationDataset(
filename_pairs,
slice_axis=slice_axis,
cache=True,
transform=tranform_lst,
slice_filter_fn=SliceFilter(
**loader_params["slice_filter_params"]))
ds.load_filenames()
if kernel_3D:
print("\nLoaded {} {} volumes of shape {}.".format(
len(ds), loader_params['slice_axis'],
context['UNet3D']['length_3D']))
else:
print("\nLoaded {} {} slices.".format(len(ds),
loader_params['slice_axis']))
# Data Loader
data_loader = DataLoader(
ds,
batch_size=context["training_parameters"]["batch_size"],
shuffle=False,
pin_memory=True,
collate_fn=imed_loader_utils.imed_collate,
num_workers=0)
# MODEL
if fname_model.endswith('.pt'):
model = torch.load(fname_model, map_location=device)
# Inference time
model.eval()
# Loop across batches
preds_list, slice_idx_list = [], []
last_sample_bool, volume, weight_matrix = False, None, None
for i_batch, batch in enumerate(data_loader):
with torch.no_grad():
img = cuda(batch['input'], cuda_available=cuda_available)
preds = model(img) if fname_model.endswith(
'.pt') else onnx_inference(fname_model, img)
preds = preds.cpu()
# Set datatype to gt since prediction should be processed the same way as gt
for modality in batch['input_metadata']:
modality[0]['data_type'] = 'gt'
# Reconstruct 3D object
for i_slice in range(len(preds)):
if "bounding_box" in batch['input_metadata'][i_slice][0]:
imed_obj_detect.adjust_undo_transforms(
undo_transforms.transforms, batch, i_slice)
if kernel_3D:
batch['gt_metadata'] = batch['input_metadata']
preds_undo, metadata, last_sample_bool, volume, weight_matrix = \
volume_reconstruction(batch, preds, undo_transforms, i_slice, volume, weight_matrix)
preds_list = [np.array(preds_undo)]
else:
# undo transformations
preds_i_undo, metadata_idx = undo_transforms(
preds[i_slice],
batch["input_metadata"][i_slice],
data_type='gt')
# Add new segmented slice to preds_list
preds_list.append(np.array(preds_i_undo))
# Store the slice index of preds_i_undo in the original 3D image
slice_idx_list.append(
int(batch['input_metadata'][i_slice][0]['slice_index']))
# If last batch and last sample of this batch, then reconstruct 3D object
if (i_batch == len(data_loader) - 1
and i_slice == len(batch['gt']) - 1) or last_sample_bool:
pred_nib = pred_to_nib(data_lst=preds_list,
fname_ref=fname_image,
fname_out=None,
z_lst=slice_idx_list,
slice_axis=slice_axis,
kernel_dim='3d' if kernel_3D else '2d',
debug=False,
bin_thr=-1)
return pred_nib
