def test_metrics_file(test_output_dirs: TestOutputDirectories) -> None:
"""Test if metrics files with Dice scores are written as expected."""
folder = test_output_dirs.make_sub_dir("test_metrics_file")
def new_file(suffix: str) -> str:
file = os.path.join(folder, suffix)
if os.path.exists(file):
os.remove(file)
return file
d = MetricsPerPatientWriter()
p1 = "Patient1"
p2 = "Patient2"
p3 = "Patient3"
liver = "liver"
kidney = "kidney"
# Ordering for test data: For "liver", patient 2 has the lowest score, sorting should move them first
# For "kidney", patient 1 has the lowest score and should be first.
d.add(p1, liver, 1.0, 1.0, 0.5)
d.add(p1, liver, 0.4, 1.0, 0.4)
d.add(p2, liver, 0.8, 1.0, 0.3)
d.add(p2, kidney, 0.7, 1.0, 0.2)
d.add(p3, kidney, 0.4, 1.0, 0.1)
metrics_file = new_file("metrics_file.csv")
d.to_csv(Path(metrics_file))
# Sorting should be first by structure name alphabetically, then Dice with lowest scores first.
assert_file_contents(
metrics_file,
"Patient,Structure,Dice,HausdorffDistance_mm,MeanDistance_mm\n"
"Patient3,kidney,0.400,1.000,0.100\n"
"Patient2,kidney,0.700,1.000,0.200\n"
"Patient1,liver,0.400,1.000,0.400\n"
"Patient2,liver,0.800,1.000,0.300\n"
"Patient1,liver,1.000,1.000,0.500\n")
aggregates_file = new_file(METRICS_AGGREGATES_FILE)
d.save_aggregates_to_csv(Path(aggregates_file))
# Sorting should be first by structure name alphabetically, then Dice with lowest scores first.
assert_file_contents_match_exactly(
Path(aggregates_file),
full_ml_test_data_path() / METRICS_AGGREGATES_FILE)
boxplot_per_structure(d.to_data_frame(),
column_name=MetricsFileColumns.DiceNumeric.value,
title="Dice score")
boxplot1 = new_file("boxplot_2class.png")
resize_and_save(5, 4, boxplot1)
plt.clf()
d.add(p1, "lung", 0.5, 2.0, 1.0)
d.add(p1, "foo", 0.9, 2.0, 1.0)
d.add(p1, "bar", 0.9, 2.0, 1.0)
d.add(p1, "baz", 0.9, 2.0, 1.0)
boxplot_per_structure(d.to_data_frame(),
column_name=MetricsFileColumns.DiceNumeric.value,
title="Dice score")
boxplot2 = new_file("boxplot_6class.png")
resize_and_save(5, 4, boxplot2)
def test_show_non_square_images(test_output_dirs: OutputFolderForTests) -> None:
input_file = full_ml_test_data_path("patch_sampling") / "scan_small.nii.gz"
input = load_nifti_image(input_file)
image = input.image
shape = image.shape
mask = np.zeros_like(image)
mask[shape[0] // 2, shape[1] // 2, shape[2] // 2] = 1
for dim in range(3):
scan_with_transparent_overlay(image, mask, dim, shape[dim] // 2, spacing=input.header.spacing)
actual_file = Path(test_output_dirs.root_dir) / f"dim_{dim}.png"
resize_and_save(5, 5, actual_file)
expected = full_ml_test_data_path("patch_sampling") / f"overlay_with_aspect_dim{dim}.png"
# To update the stored results, uncomment this line:
# expected.write_bytes(actual_file.read_bytes())
assert_binary_files_match(actual_file, expected)
def test_plot_dice_per_epoch(test_output_dirs: OutputFolderForTests, num_classes: int) -> None:
metrics: Dict[str, Any] = {}
epoch = [1, 2, 3]
for i in range(num_classes):
metric_name = "Val_Dice/Class{}".format(i)
loss = [i / num_classes * j / 3 for j in epoch]
metrics[metric_name] = {"epoch": epoch, "loss": loss}
metrics["baz"] = [17]
series_count = plotting.plot_val_dice_per_epoch(metrics)
file_name = test_output_dirs.root_dir / f"dice_per_epoch_{num_classes}classes.png"
plotting.add_legend(series_count)
plotting.resize_and_save(5, 4, file_name)
assert file_name.is_file()
# Try writing the same figure again, to see what the file overwrite behaviour is.
# In actual training runs, the file will be overwritten repeatedly.
plotting.resize_and_save(5, 4, file_name)
def test_plot_overlay(test_output_dirs: OutputFolderForTests,
dimension: int) -> None:
set_random_seed(0)
shape = (10, 30, 30)
image = np.random.rand(*shape).astype(np.float32) * 1000
mask = np.zeros(shape).flatten()
for i in range(len(mask)):
mask[i] = i
mask = mask.reshape(shape)
plt.figure()
scan_with_transparent_overlay(image, mask, dimension, shape[dimension] // 2, spacing=(1.0, 1.0, 1.0))
file = Path(test_output_dirs.root_dir) / "plot.png"
resize_and_save(5, 5, file)
assert file.exists()
expected = full_ml_test_data_path("patch_sampling") / f"overlay_{dimension}.png"
# To update the stored results, uncomment this line:
# expected.write_bytes(file.read_bytes())
assert_binary_files_match(file, expected)
def segmentation_model_test_epoch(
config: SegmentationModelBase,
data_split: ModelExecutionMode,
test_epoch: int,
results_folder: Path,
epoch_and_split: str,
run_recovery: Optional[RunRecovery] = None) -> Optional[List[float]]:
"""
The main testing loop for a given epoch. It loads the model and datasets, then proceeds to test the model.
Returns a list with an entry for each image in the dataset. The entry is the average Dice score,
where the average is taken across all non-background structures in the image.
:param test_epoch: The last trained epoch of the model.
:param config: The arguments which specify all required information.
:param data_split: Is the model evaluated on train, test, or validation set?
:param results_folder: The folder where to store the results
:param epoch_and_split: A string that should uniquely identify the epoch and the data split (train/val/test).
:param run_recovery: Run recovery data if applicable.
:raises TypeError: If the arguments are of the wrong type.
:raises ValueError: When there are issues loading the model.
:return A list with the mean dice score (across all structures apart from background) for each image.
"""
ml_util.set_random_seed(config.get_effective_random_seed(),
"Model Training")
results_folder = Path(results_folder)
results_folder.mkdir(exist_ok=True)
test_dataframe = config.get_dataset_splits()[data_split]
test_csv_path = results_folder / STORED_CSV_FILE_NAMES[data_split]
test_dataframe.to_csv(path_or_buf=test_csv_path, index=False)
logging.info("Results directory: {}".format(results_folder))
logging.info(
f"Starting evaluation of model {config.model_name} on {epoch_and_split}"
)
# Write the dataset id and ground truth ids into the results folder
store_run_information(results_folder, config.azure_dataset_id,
config.ground_truth_ids, config.image_channels)
ds = config.get_torch_dataset_for_inference(data_split)
inference_pipeline = create_inference_pipeline(config=config,
epoch=test_epoch,
run_recovery=run_recovery)
if inference_pipeline is None:
# This will happen if there is no checkpoint for the given epoch, in either the recovered run (if any) or
# the current one.
return None
# for mypy
assert isinstance(inference_pipeline, FullImageInferencePipelineBase)
# Deploy the trained model on a set of images and store output arrays.
for sample_index, sample in enumerate(ds, 1):
logging.info(f"Predicting for image {sample_index} of {len(ds)}...")
sample = Sample.from_dict(sample=sample)
inference_result = inference_pipeline.predict_and_post_process_whole_image(
image_channels=sample.image,
mask=sample.mask,
patient_id=sample.patient_id,
voxel_spacing_mm=sample.metadata.image_header.spacing)
store_inference_results(inference_result=inference_result,
config=config,
results_folder=results_folder,
image_header=sample.metadata.image_header)
# Evaluate model generated segmentation maps.
num_workers = min(cpu_count(), len(ds))
with Pool(processes=num_workers) as pool:
pool_outputs = pool.map(
partial(evaluate_model_predictions,
config=config,
dataset=ds,
results_folder=results_folder), range(len(ds)))
average_dice = list()
metrics_writer = MetricsPerPatientWriter()
for (patient_metadata, metrics_for_patient) in pool_outputs:
# Add the Dice score for the foreground classes, stored in the default hue
metrics.add_average_foreground_dice(metrics_for_patient)
average_dice.append(
metrics_for_patient.get_single_metric(MetricType.DICE))
# Structure names does not include the background class (index 0)
for structure_name in config.ground_truth_ids:
dice_for_struct = metrics_for_patient.get_single_metric(
MetricType.DICE, hue=structure_name)
hd_for_struct = metrics_for_patient.get_single_metric(
MetricType.HAUSDORFF_mm, hue=structure_name)
md_for_struct = metrics_for_patient.get_single_metric(
MetricType.MEAN_SURFACE_DIST_mm, hue=structure_name)
metrics_writer.add(patient=str(patient_metadata.patient_id),
structure=structure_name,
dice=dice_for_struct,
hausdorff_distance_mm=hd_for_struct,
mean_distance_mm=md_for_struct)
metrics_writer.to_csv(results_folder / METRICS_FILE_NAME)
metrics_writer.save_aggregates_to_csv(results_folder /
METRICS_AGGREGATES_FILE)
if config.is_plotting_enabled:
plt.figure()
boxplot_per_structure(metrics_writer.to_data_frame(),
column_name=MetricsFileColumns.DiceNumeric.value,
title=f"Dice score for {epoch_and_split}")
# The box plot file will be written to the output directory. AzureML will pick that up, and display
# on the run overview page, without having to log to the run context.
plotting.resize_and_save(5, 4, results_folder / BOXPLOT_FILE)
plt.close()
logging.info(
f"Finished evaluation of model {config.model_name} on {epoch_and_split}"
)
return average_dice
def visualize_random_crops(sample: Sample, config: SegmentationModelBase,
output_folder: Path) -> np.ndarray:
"""
Simulate the effect of sampling random crops (as is done for trainig segmentation models), and store the results
as a Nifti heatmap and as 3 axial/sagittal/coronal slices. The heatmap and the slices are stored in the given
output folder, with filenames that contain the patient ID as the prefix.
:param sample: The patient information from the dataset, with scans and ground truth labels.
:param config: The model configuration.
:param output_folder: The folder into which the heatmap and thumbnails should be written.
:return: A numpy array that has the same size as the image, containing how often each voxel was contained in
"""
output_folder.mkdir(exist_ok=True, parents=True)
sample = CroppingDataset.create_possibly_padded_sample_for_cropping(
sample=sample,
crop_size=config.crop_size,
padding_mode=config.padding_mode)
logging.info(f"Processing sample: {sample.patient_id}")
# Exhaustively sample with random crop function
image_channel0 = sample.image[0]
heatmap = np.zeros(image_channel0.shape, dtype=np.uint16)
# Number of repeats should fit into the range of UInt16, because we will later save the heatmap as an integer
# Nifti file of that datatype.
repeats = 200
for _ in range(repeats):
slicers, _ = augmentation.slicers_for_random_crop(
sample=sample,
crop_size=config.crop_size,
class_weights=config.class_weights)
heatmap[slicers[0], slicers[1], slicers[2]] += 1
is_3dim = heatmap.shape[0] > 1
header = sample.metadata.image_header
if not header:
logging.warning(
f"No image header found for patient {sample.patient_id}. Using default header."
)
header = get_unit_image_header()
if is_3dim:
ct_output_name = str(output_folder / f"{sample.patient_id}_ct.nii.gz")
heatmap_output_name = str(
output_folder / f"{sample.patient_id}_sampled_patches.nii.gz")
io_util.store_as_nifti(image=heatmap,
header=header,
file_name=heatmap_output_name,
image_type=heatmap.dtype,
scale=False)
io_util.store_as_nifti(image=image_channel0,
header=header,
file_name=ct_output_name,
image_type=sample.image.dtype,
scale=False)
heatmap_scaled = heatmap.astype(dtype=np.float) / heatmap.max()
# If the incoming image is effectively a 2D image with degenerate Z dimension, then only plot a single
# axial thumbnail. Otherwise, plot thumbnails for all 3 dimensions.
dimensions = list(range(3)) if is_3dim else [0]
# Center the 3 thumbnails at one of the points where the heatmap attains a maximum. This should ensure that
# the thumbnails are in an area where many of the organs of interest are located.
max_heatmap_index = np.unravel_index(
heatmap.argmax(), heatmap.shape) if is_3dim else (0, 0, 0)
for dimension in dimensions:
plt.clf()
scan_with_transparent_overlay(
scan=image_channel0,
overlay=heatmap_scaled,
dimension=dimension,
position=max_heatmap_index[dimension] if is_3dim else 0,
spacing=header.spacing)
# Construct a filename that has a dimension suffix if we are generating 3 of them. For 2dim images, skip
# the suffix.
thumbnail = f"{sample.patient_id}_sampled_patches"
if is_3dim:
thumbnail += f"_dim{dimension}"
thumbnail += ".png"
resize_and_save(width_inch=5,
height_inch=5,
filename=output_folder / thumbnail)
return heatmap
