def test_anomaly_detection(value_to_insert: float,
in_training_mode: bool) -> None:
"""
Test anomaly detection for the segmentation forward pass.
:param value_to_insert: The value to insert in the image image (nan, inf, or a valid float)
:param in_training_mode: If true, run the segmentation forward pass in training mode, otherwise use the
settings for running on the validation set.
:return:
"""
image_size = [1, 1, 4, 4, 4]
labels_size = [1, 2, 4, 4, 4]
mask_size = [1, 4, 4, 4]
crop_size = (4, 4, 4)
inference_stride_size = (2, 2, 2)
ground_truth_ids = ["Lung"]
# image to run inference on
image = torch.from_numpy(
np.random.uniform(size=image_size).astype(ImageDataType.IMAGE.value))
# labels for criterion
labels = torch.from_numpy(
np.random.uniform(size=labels_size).astype(
ImageDataType.SEGMENTATION.value))
# create a random mask if required
mask = torch.from_numpy((np.round(np.random.uniform(
size=mask_size)).astype(dtype=ImageDataType.MASK.value)))
config = SegmentationModelBase(crop_size=crop_size,
inference_stride_size=inference_stride_size,
image_channels=["ct"],
ground_truth_ids=ground_truth_ids,
should_validate=False,
detect_anomaly=True)
# instantiate the model
model = SimpleModel(1, [1], 2, 2)
config.adjust_after_mixed_precision_and_parallel(model)
config.use_gpu = False
# Create the optimizer_type and loss criterion
optimizer = model_util.create_optimizer(config, model)
criterion = lambda x, y: torch.tensor(value_to_insert, requires_grad=True)
pipeline = SegmentationForwardPass(model,
config,
batch_size=1,
optimizer=optimizer,
in_training_mode=in_training_mode,
criterion=criterion)
image[0, 0, 0, 0, 0] = value_to_insert
if np.isnan(value_to_insert) or np.isinf(value_to_insert):
with pytest.raises(RuntimeError) as ex:
pipeline.forward_pass_patches(patches=image,
mask=mask,
labels=labels)
assert f"loss computation returned {value_to_insert}" in str(ex)
else:
pipeline.forward_pass_patches(patches=image, mask=mask, labels=labels)
def test_dataset_csv_with_SegmentationModelBase(
test_output_dirs: OutputFolderForTests) -> None:
dataset_csv_path = create_dataset_csv(test_output_dirs)
model_config = SegmentationModelBase(should_validate=False)
model_config.local_dataset = dataset_csv_path.parent
model_config.dataset_csv = dataset_csv_path.name
dataframe = model_config.read_dataset_if_needed()
assert dataframe is not None
validate_dataset_paths(model_config)
def random_select_patch_center(sample: Sample,
class_weights: List[float] = None
) -> np.ndarray:
"""
Samples a point to use as the coordinates of the patch center. First samples one
class among the available classes then samples a center point among the pixels of the sampled
class.
:param sample: A set of Image channels, ground truth labels and mask to randomly crop.
:param class_weights: A weighting vector with values [0, 1] to influence the class the center crop
voxel belongs to (must sum to 1), uniform distribution assumed if none provided.
:return numpy int array (3x1) containing patch center spatial coordinates
"""
num_classes = sample.labels.shape[0]
if class_weights is not None:
if len(class_weights) != num_classes:
raise Exception(
"A weight must be provided for each class, found weights:{}, expected:{}"
.format(len(class_weights), num_classes))
SegmentationModelBase.validate_class_weights(class_weights)
# If class weights are not initialised, selection is made with equal probability for all classes
available_classes = list(range(num_classes))
original_class_weights = class_weights
while len(available_classes) > 0:
selected_label_class = random.choices(population=available_classes,
weights=class_weights,
k=1)[0]
# Check pixels where mask and label maps are both foreground
indices = np.argwhere(
np.logical_and(sample.labels[selected_label_class] == 1.0,
sample.mask == 1))
if not np.any(indices):
available_classes.remove(selected_label_class)
if class_weights is not None:
assert original_class_weights is not None # for mypy
class_weights = [
original_class_weights[i] for i in available_classes
]
if sum(class_weights) <= 0.0:
raise ValueError(
"Cannot sample a class: no class present in the sample has a positive weight"
)
else:
break
# Raise an exception if non of the foreground classes are overlapping with the mask
if len(available_classes) == 0:
raise Exception(
"No non-mask voxels found, please check your mask and labels map")
# noinspection PyUnboundLocalVariable
choice = random.randint(0, len(indices) - 1)
return indices[choice].astype(int) # Numpy usually stores as floats
def segmentation_model_test(
config: SegmentationModelBase,
data_split: ModelExecutionMode,
checkpoint_handler: CheckpointHandler,
model_proc: ModelProcessing = ModelProcessing.DEFAULT
) -> InferenceMetricsForSegmentation:
"""
The main testing loop for segmentation models.
It loads the model and datasets, then proceeds to test the model for all requested checkpoints.
:param config: The arguments object which has a valid random seed attribute.
:param data_split: Indicates which of the 3 sets (training, test, or validation) is being processed.
:param checkpoint_handler: Checkpoint handler object to find checkpoint paths for model initialization
:param model_proc: whether we are testing an ensemble or single model
:return: InferenceMetric object that contains metrics related for all of the checkpoint epochs.
"""
results: Dict[int, float] = {}
checkpoints_to_test = checkpoint_handler.get_checkpoints_to_test()
if not checkpoints_to_test:
raise ValueError(
"There were no checkpoints available for model testing.")
for checkpoint_paths_and_epoch in checkpoints_to_test:
epoch = checkpoint_paths_and_epoch.epoch
epoch_results_folder = config.outputs_folder / get_epoch_results_path(
epoch, data_split, model_proc)
# save the datasets.csv used
config.write_dataset_files(root=epoch_results_folder)
epoch_and_split = "epoch {} {} set".format(epoch, data_split.value)
epoch_dice_per_image = segmentation_model_test_epoch(
config=copy.deepcopy(config),
data_split=data_split,
checkpoint_paths=checkpoint_paths_and_epoch.checkpoint_paths,
results_folder=epoch_results_folder,
epoch_and_split=epoch_and_split)
if epoch_dice_per_image is None:
logging.warning(
"There is no checkpoint file for epoch {}".format(epoch))
else:
epoch_average_dice: float = np.mean(
epoch_dice_per_image) if len(epoch_dice_per_image) > 0 else 0
results[epoch] = epoch_average_dice
logging.info("Epoch: {:3} | Mean Dice: {:4f}".format(
epoch, epoch_average_dice))
if model_proc == ModelProcessing.ENSEMBLE_CREATION:
# For the upload, we want the path without the "OTHER_RUNS/ENSEMBLE" prefix.
name = str(
get_epoch_results_path(epoch, data_split,
ModelProcessing.DEFAULT))
PARENT_RUN_CONTEXT.upload_folder(
name=name, path=str(epoch_results_folder))
if len(results) == 0:
raise ValueError(
"There was no single checkpoint file available for model testing.")
return InferenceMetricsForSegmentation(data_split=data_split,
epochs=results)
def test_set_model_config_attributes() -> None:
"""Tests setter function for model config attributes"""
train_output_size = (3, 5, 3)
test_output_size = (7, 7, 7)
model = IdentityModel()
model_config = SegmentationModelBase(crop_size=train_output_size,
test_crop_size=test_output_size,
should_validate=False)
model_config.adjust_after_mixed_precision_and_parallel(model)
assert model_config.inference_stride_size == test_output_size
def test_non_overridable_properties() -> None:
"""
Test to make sure properties that are private or marked as constant/readonly/NON OVERRIDABLE are not
configurable through the commandline.
"""
non_overridable = ["--" + k + "=" + str(v.default) for k, v in SegmentationModelBase.params().items()
if k not in SegmentationModelBase.get_overridable_parameters().keys()]
parser = SegmentationModelBase.create_argparser()
# try to parse the non overridable arguments
_, unknown = parser.parse_known_args(non_overridable)
assert all([x in unknown for x in non_overridable])
def test_set_model_config_attributes() -> None:
"""Tests setter function for model config attributes"""
train_output_size = (3, 5, 3)
test_output_size = (7, 7, 7)
model = IdentityModel()
model_config = SegmentationModelBase(crop_size=train_output_size,
test_crop_size=test_output_size,
should_validate=False)
model_config.set_derived_model_properties(model)
assert model_config.inference_stride_size == test_output_size
def test_download_checkpoints_hyperdrive_run(test_output_dirs: TestOutputDirectories,
runner_config: AzureConfig) -> None:
output_dir = Path(test_output_dirs.root_dir)
config = SegmentationModelBase(should_validate=False)
config.set_output_to(output_dir)
runner_config.run_recovery_id = DEFAULT_ENSEMBLE_RUN_RECOVERY_ID
child_runs = fetch_child_runs(run=fetch_run(runner_config.get_workspace(), DEFAULT_ENSEMBLE_RUN_RECOVERY_ID))
# recover child runs separately also to test hyperdrive child run recovery functionality
expected_checkpoint_file = "1" + CHECKPOINT_FILE_SUFFIX
for child in child_runs:
expected_files = [Path(config.checkpoint_folder) / child.id / expected_checkpoint_file]
run_recovery = RunRecovery.download_checkpoints_from_recovery_run(runner_config, config, child)
assert all([x in expected_files for x in run_recovery.get_checkpoint_paths(epoch=1)])
assert all([expected_file.exists() for expected_file in expected_files])
def test_save_dataset_example(test_output_dirs: OutputFolderForTests) -> None:
"""
Test if the example dataset can be saved as expected.
"""
image_size = (10, 20, 30)
label_size = (2, ) + image_size
spacing = (1, 2, 3)
np.random.seed(0)
# Image should look similar to what a photonormalized image looks like: Centered around 0
image = np.random.rand(*image_size) * 2 - 1
# Labels are expected in one-hot encoding, predictions as class index
labels = np.zeros(label_size, dtype=int)
labels[0] = 1
labels[0, 5:6, 10:11, 15:16] = 0
labels[1, 5:6, 10:11, 15:16] = 1
prediction = np.zeros(image_size, dtype=int)
prediction[4:7, 9:12, 14:17] = 1
dataset_sample = DatasetExample(epoch=1,
patient_id=2,
header=ImageHeader(origin=(0, 1, 0),
direction=(1, 0, 0, 0,
1, 0, 0, 0,
1),
spacing=spacing),
image=image,
prediction=prediction,
labels=labels)
images_folder = test_output_dirs.root_dir
config = SegmentationModelBase(
should_validate=False,
norm_method=PhotometricNormalizationMethod.Unchanged)
config.set_output_to(images_folder)
store_and_upload_example(dataset_sample, config)
image_from_disk = io_util.load_nifti_image(
os.path.join(config.example_images_folder, "p2_e_1_image.nii.gz"))
labels_from_disk = io_util.load_nifti_image(
os.path.join(config.example_images_folder, "p2_e_1_label.nii.gz"))
prediction_from_disk = io_util.load_nifti_image(
os.path.join(config.example_images_folder, "p2_e_1_prediction.nii.gz"))
assert image_from_disk.header.spacing == spacing
# When no photometric normalization is provided when saving, image is multiplied by 1000.
# It is then rounded to int64, but converted back to float when read back in.
expected_from_disk = (image * 1000).astype(np.int16).astype(np.float64)
assert np.array_equal(image_from_disk.image, expected_from_disk)
assert labels_from_disk.header.spacing == spacing
assert np.array_equal(labels_from_disk.image, np.argmax(labels, axis=0))
assert prediction_from_disk.header.spacing == spacing
assert np.array_equal(prediction_from_disk.image, prediction)
def test_store_image_as_short_nifti(test_output_dirs: TestOutputDirectories,
norm_method: PhotometricNormalizationMethod,
image_range: Any,
window_level: Any) -> None:
window, level = window_level if window_level else (400, 0)
image = np.random.random_sample((1, 2, 3))
image_shape = image.shape
args = SegmentationModelBase(norm_method=norm_method, window=window, level=level, should_validate=False)
# Get integer values that are in the image range
image1 = LinearTransform.transform(data=image, input_range=(0, 1), output_range=args.output_range)
image = image1.astype(np.short) # type: ignore
header = ImageHeader(origin=(1, 1, 1), direction=(1, 0, 0, 0, 1, 0, 0, 0, 1), spacing=(1, 1, 1))
nifti_name = test_output_dirs.create_file_or_folder_path(default_image_name)
io_util.store_image_as_short_nifti(image, header, nifti_name, args)
if norm_method == PhotometricNormalizationMethod.CtWindow:
output_range = get_range_for_window_level(args.level, args.window)
image = LinearTransform.transform(data=image, input_range=args.output_range, output_range=output_range)
image = image.astype(np.short)
else:
image = image * 1000
t = np.unique(image)
assert_nifti_content(nifti_name, image_shape, header, list(t), np.short)
def create_inference_pipeline(model_config: SegmentationModelBase,
full_path_to_checkpoints: List[Path],
use_gpu: bool = True) \
-> Tuple[FullImageInferencePipelineBase, SegmentationModelBase]:
"""
Create pipeline for inference, this can be a single model inference pipeline or an ensemble, if multiple
checkpoints provided.
:param model_config: Model config to use to create the pipeline.
:param full_path_to_checkpoints: Checkpoints to use for model inference.
:param use_gpu: If GPU should be used or not.
"""
model_config.use_gpu = use_gpu
logging.info('test_config: ' + model_config.model_name)
inference_pipeline: Optional[FullImageInferencePipelineBase]
if len(full_path_to_checkpoints) == 1:
inference_pipeline = InferencePipeline.create_from_checkpoint(
path_to_checkpoint=full_path_to_checkpoints[0],
model_config=model_config)
else:
inference_pipeline = EnsemblePipeline.create_from_checkpoints(
path_to_checkpoints=full_path_to_checkpoints,
model_config=model_config)
if inference_pipeline is None:
raise ValueError("Cannot create inference pipeline")
return inference_pipeline, model_config
def test_plot_normalization_result(
test_output_dirs: TestOutputDirectories) -> None:
"""
Tests plotting of before/after histograms in photometric normalization.
:return:
"""
size = (3, 3, 3)
image = np.zeros((1, ) + size)
for i, (z, y, x) in enumerate(
itertools.product(range(size[0]), range(size[1]), range(size[2]))):
image[0, z, y, x] = i
labels = np.zeros((2, ) + size)
labels[1, 1, 1, 1] = 1
sample = Sample(image=image,
labels=labels,
mask=np.ones(size),
metadata=DummyPatientMetadata)
config = SegmentationModelBase(
norm_method=PhotometricNormalizationMethod.CtWindow,
window=4,
level=13,
should_validate=False)
normalizer = PhotometricNormalization(config)
folder = Path(test_output_dirs.root_dir)
files = plotting.plot_normalization_result(sample, normalizer, folder)
expected = ["042_slice_001.png", "042_slice_001_contour.png"]
compare_files(files, expected)
def full_image_dataset(default_config: SegmentationModelBase,
normalize_fn: Callable) -> FullImageDataset:
df = default_config.get_dataset_splits()
return FullImageDataset(
args=default_config,
full_image_sample_transforms=Compose3D([normalize_fn]), # type: ignore
data_frame=df.train)
def test_validate_inference_stride_size() -> None:
SegmentationModelBase.validate_inference_stride_size(
inference_stride_size=(5, 3, 1), output_size=(5, 3, 1))
SegmentationModelBase.validate_inference_stride_size(
inference_stride_size=(5, 3, 1), output_size=None)
with pytest.raises(ValueError):
SegmentationModelBase.validate_inference_stride_size(
inference_stride_size=(5, 3, 1), output_size=(3, 3, 3))
SegmentationModelBase.validate_inference_stride_size(
inference_stride_size=(5, 3, 0), output_size=None)
def segmentation_model_test(
config: SegmentationModelBase,
execution_mode: ModelExecutionMode,
checkpoint_paths: List[Path],
model_proc: ModelProcessing = ModelProcessing.DEFAULT
) -> InferenceMetricsForSegmentation:
"""
The main testing loop for segmentation models.
It loads the model and datasets, then proceeds to test the model for all requested checkpoints.
:param config: The arguments object which has a valid random seed attribute.
:param execution_mode: Indicates which of the 3 sets (training, test, or validation) is being processed.
:param checkpoint_handler: Checkpoint handler object to find checkpoint paths for model initialization.
:param model_proc: Whether we are testing an ensemble or single model.
:param patient_id: String which contains subject identifier.
:return: InferenceMetric object that contains metrics related for all of the checkpoint epochs.
"""
epoch_results_folder = config.outputs_folder / get_best_epoch_results_path(
execution_mode, model_proc)
# save the datasets.csv used
config.write_dataset_files(root=epoch_results_folder)
epoch_and_split = f"{execution_mode.value} set"
epoch_dice_per_image = segmentation_model_test_epoch(
config=copy.deepcopy(config),
execution_mode=execution_mode,
checkpoint_paths=checkpoint_paths,
results_folder=epoch_results_folder,
epoch_and_split=epoch_and_split)
if epoch_dice_per_image is None:
raise ValueError(
"There was no single checkpoint file available for model testing.")
else:
epoch_average_dice: float = np.mean(
epoch_dice_per_image) if len(epoch_dice_per_image) > 0 else 0
result = epoch_average_dice
logging.info(f"Mean Dice: {epoch_average_dice:4f}")
if model_proc == ModelProcessing.ENSEMBLE_CREATION:
# For the upload, we want the path without the "OTHER_RUNS/ENSEMBLE" prefix.
name = str(
get_best_epoch_results_path(execution_mode,
ModelProcessing.DEFAULT))
PARENT_RUN_CONTEXT.upload_folder(name=name,
path=str(epoch_results_folder))
return InferenceMetricsForSegmentation(execution_mode=execution_mode,
metrics=result)
def main(yaml_file_path: Path) -> None:
"""
Invoke either by
* specifying a model, '--model Lung'
* or specifying dataset and normalization parameters separately: --azure_dataset_id=foo --norm_method=None
In addition, the arguments '--image_channel' and '--gt_channel' must be specified (see below).
"""
config, runner_config, args = get_configs(
SegmentationModelBase(should_validate=False), yaml_file_path)
dataset_config = DatasetConfig(name=config.azure_dataset_id,
local_folder=config.local_dataset,
use_mounting=True)
local_dataset, mount_context = dataset_config.to_input_dataset_local(
workspace=runner_config.get_workspace())
dataframe = pd.read_csv(local_dataset / DATASET_CSV_FILE_NAME)
normalizer_config = NormalizeAndVisualizeConfig(**args)
actual_mask_channel = None if normalizer_config.ignore_mask else config.mask_id
image_channel = normalizer_config.image_channel or config.image_channels[0]
if not image_channel:
raise ValueError(
"No image channel selected. Specify a model by name, or use the image_channel argument."
)
gt_channel = normalizer_config.gt_channel or config.ground_truth_ids[0]
if not gt_channel:
raise ValueError(
"No GT channel selected. Specify a model by name, or use the gt_channel argument."
)
dataset_sources = load_dataset_sources(
dataframe,
local_dataset_root_folder=local_dataset,
image_channels=[image_channel],
ground_truth_channels=[gt_channel],
mask_channel=actual_mask_channel)
result_folder = local_dataset
if normalizer_config.result_folder is not None:
result_folder = result_folder / normalizer_config.result_folder
if not result_folder.is_dir():
result_folder.mkdir()
all_patient_ids = [*dataset_sources.keys()]
if normalizer_config.only_first == 0:
patient_ids_to_process = all_patient_ids
else:
patient_ids_to_process = all_patient_ids[:normalizer_config.only_first]
args_file = result_folder / ARGS_TXT
args_file.write_text(" ".join(sys.argv[1:]))
config_file = result_folder / "config.txt"
config_file.write_text(str(config))
normalizer = PhotometricNormalization(config)
for patient_id in patient_ids_to_process:
print(f"Starting to process patient {patient_id}")
images = load_images_from_dataset_source(dataset_sources[patient_id])
plotting.plot_normalization_result(images,
normalizer,
result_folder,
result_prefix=image_channel)
def test_crop_size_multiple_in_build_net() -> None:
"""
Tests if the the crop_size validation is really called in the model creation code
"""
config = SegmentationModelBase(architecture=ModelArchitectureConfig.UNet3D,
image_channels=["ct"],
feature_channels=[1],
kernel_size=3,
crop_size=(17, 16, 16),
should_validate=False)
# Crop size of 17 in dimension 0 is invalid for a UNet3D, should be multiple of 16.
# This should raise a ValueError.
with pytest.raises(ValueError) as ex:
build_net(config)
assert "Training crop size: Crop size is not valid" in str(ex)
config.crop_size = (16, 16, 16)
config.test_crop_size = (17, 18, 19)
with pytest.raises(ValueError) as ex:
build_net(config)
assert "Test crop size: Crop size is not valid" in str(ex)
def test_fields_are_set() -> None:
"""
Tests that expected fields are set when creating config classes.
"""
expected = [("hello", None), ("world", None)]
config = SegmentationModelBase(
should_validate=False,
ground_truth_ids=[x[0] for x in expected],
largest_connected_component_foreground_classes=expected)
assert hasattr(config, CROSS_VALIDATION_SPLIT_INDEX_TAG_KEY)
assert config.largest_connected_component_foreground_classes == expected
def test_inference_required_crossval_runs() -> None:
"""
Test the flags for running full inference on the test set, for models that are trained in crossval mode.
"""
classification_model = GlaucomaPublic()
classification_model.number_of_cross_validation_splits = 2
segmentation_model = SegmentationModelBase(should_validate=False)
segmentation_model.number_of_cross_validation_splits = 2
assert classification_model.perform_cross_validation
assert segmentation_model.perform_cross_validation
# Cross validation child runs for classification models need test set inference to ensure that the report works
# correctly.
assert classification_model.is_inference_required(
model_proc=ModelProcessing.DEFAULT, data_split=ModelExecutionMode.TEST)
# For models other than classification models, there is by default no inference on the test set.
assert not segmentation_model.is_inference_required(
model_proc=ModelProcessing.DEFAULT, data_split=ModelExecutionMode.TEST)
classification_model.inference_on_test_set = False
# If a flag is set explicitly, use that.
assert not classification_model.is_inference_required(
model_proc=ModelProcessing.DEFAULT, data_split=ModelExecutionMode.TEST)
def compare_scores_against_baselines(model_config: SegmentationModelBase,
azure_config: AzureConfig,
model_proc: ModelProcessing) -> None:
"""
If the model config has any baselines to compare against, loads the metrics.csv file that should just have
been written for the last epoch of the current run, and its dataset.csv. Do the same for all the baselines,
whose corresponding files should be in the repository already. For each baseline, call the Wilcoxon signed-rank test
on pairs consisting of Dice scores from the current model and the baseline, and print out comparisons to
the Wilcoxon results file.
"""
# The attribute will only be present for a segmentation model; and it might be None or empty even for that.
comparison_blob_storage_paths = model_config.comparison_blob_storage_paths
if not comparison_blob_storage_paths:
return
outputs_path = model_config.outputs_folder / get_best_epoch_results_path(
ModelExecutionMode.TEST, model_proc)
if not outputs_path.is_dir():
if not model_config.is_inference_required(model_proc,
ModelExecutionMode.TEST):
logging.info(INFERENCE_DISABLED_WARNING)
return
raise FileNotFoundError(
f"Cannot compare scores against baselines: no best epoch results found at {outputs_path}"
)
model_metrics_path = outputs_path / SUBJECT_METRICS_FILE_NAME
model_dataset_path = outputs_path / DATASET_CSV_FILE_NAME
if not model_dataset_path.exists():
raise FileNotFoundError(
f"Not comparing with baselines because no {model_dataset_path} file found for this run"
)
if not model_metrics_path.exists():
raise FileNotFoundError(
f"Not comparing with baselines because no {model_metrics_path} file found for this run"
)
model_metrics_df = pd.read_csv(model_metrics_path)
model_dataset_df = pd.read_csv(model_dataset_path)
comparison_result = download_and_compare_scores(
outputs_path, azure_config, comparison_blob_storage_paths,
model_dataset_df, model_metrics_df)
full_metrics_path = str(outputs_path / FULL_METRICS_DATAFRAME_FILE)
comparison_result.dataframe.to_csv(full_metrics_path)
if comparison_result.did_comparisons:
wilcoxon_path = outputs_path / BASELINE_WILCOXON_RESULTS_FILE
logging.info(
f"Wilcoxon tests of current {model_proc.value} model against baseline(s), "
f"written to {wilcoxon_path}:")
for line in comparison_result.wilcoxon_lines:
logging.info(line)
logging.info("End of Wilcoxon test results")
may_write_lines_to_file(comparison_result.wilcoxon_lines,
wilcoxon_path)
write_to_scatterplot_directory(outputs_path, comparison_result.plots)
def __init__(self, config: SegmentationModelBase, *args: Any, **kwargs: Any) -> None:
super().__init__(config, *args, **kwargs)
self.model = config.create_model()
self.loss_fn = model_util.create_segmentation_loss_function(config)
self.ground_truth_ids = config.ground_truth_ids
self.train_dice = MetricForMultipleStructures(ground_truth_ids=self.ground_truth_ids, is_training=True)
self.val_dice = MetricForMultipleStructures(ground_truth_ids=self.ground_truth_ids, is_training=False)
self.train_voxels = MetricForMultipleStructures(ground_truth_ids=self.ground_truth_ids, is_training=True,
metric_name=MetricType.VOXEL_COUNT.value,
use_average_across_structures=False)
self.val_voxels = MetricForMultipleStructures(ground_truth_ids=self.ground_truth_ids, is_training=False,
metric_name=MetricType.VOXEL_COUNT.value,
use_average_across_structures=False)
def test_get_output_size() -> None:
"""Tests config properties related to output tensor size"""
train_output_size = (5, 5, 5)
test_output_size = (7, 7, 7)
model_config = SegmentationModelBase(crop_size=train_output_size,
test_crop_size=test_output_size,
should_validate=False)
assert model_config.get_output_size(execution_mode=ModelExecutionMode.TRAIN) is None
assert model_config.get_output_size(execution_mode=ModelExecutionMode.TEST) is None
model = IdentityModel()
model_config.set_derived_model_properties(model)
assert model_config.get_output_size(execution_mode=ModelExecutionMode.TRAIN) == train_output_size
assert model_config.get_output_size(execution_mode=ModelExecutionMode.TEST) == test_output_size
def test_invalid_stride_size(test_output_dirs: OutputFolderForTests) -> None:
config = SegmentationModelBase(
architecture="UNet3D",
feature_channels=[1],
crop_size=(64, 64, 64),
test_crop_size=(80, 80, 80),
image_channels=["mr"],
ground_truth_ids=["tumour_mass", "subtract"],
train_batch_size=8,
inference_batch_size=1,
inference_stride_size=(120, 120, 120),
should_validate=False
)
config.set_output_to(test_output_dirs.root_dir)
checkpoint_path = test_output_dirs.root_dir / "checkpoint.ckpt"
create_model_and_store_checkpoint(config, checkpoint_path)
with pytest.raises(ValueError) as ex:
load_from_checkpoint_and_adjust_for_inference(config=config, checkpoint_path=checkpoint_path)
assert "The inference stride size (120, 120, 120) must be smaller" in ex.value.args[0]
assert str(config.inference_stride_size) in ex.value.args[0]
assert str(config.test_crop_size) in ex.value.args[0]
def test_inference_stride_size_setter() -> None:
"""Tests setter function raises an error when stride size is larger than output patch size"""
test_output_size = (7, 3, 5)
test_stride_size = (3, 3, 3)
test_fail_stride_size = (1, 1, 9)
model = IdentityModel()
model_config = SegmentationModelBase(test_crop_size=test_output_size, should_validate=False)
model_config.inference_stride_size = test_stride_size
assert model_config.inference_stride_size == test_stride_size
model_config.set_derived_model_properties(model)
assert model_config.inference_stride_size == test_stride_size
model_config.inference_stride_size = None
model_config.set_derived_model_properties(model)
assert model_config.inference_stride_size == test_output_size
with pytest.raises(ValueError):
model_config.inference_stride_size = test_fail_stride_size
def test_copy_child_paths_to_folder(
is_ensemble: bool, extra_code_directory: str,
test_output_dirs: OutputFolderForTests) -> None:
azure_config = AzureConfig(extra_code_directory=extra_code_directory)
fake_model = SegmentationModelBase(should_validate=False)
fake_model.set_output_to(test_output_dirs.root_dir)
# To simulate ensemble models, there are two checkpoints, one in the root dir and one in a folder
checkpoints_absolute, checkpoints_relative = create_checkpoints(
fake_model, is_ensemble)
# Simulate a project root: We can't derive that from the repository root because that might point
# into Python's package folder
project_root = Path(__file__).parent.parent
ml_runner = MLRunner(model_config=fake_model,
azure_config=azure_config,
project_root=project_root)
model_folder = test_output_dirs.root_dir / "final"
ml_runner.copy_child_paths_to_folder(model_folder=model_folder,
checkpoint_paths=checkpoints_absolute)
expected_files = [
fixed_paths.ENVIRONMENT_YAML_FILE_NAME,
fixed_paths.MODEL_INFERENCE_JSON_FILE_NAME,
"InnerEye/ML/runner.py",
"InnerEye/ML/model_testing.py",
"InnerEye/Common/fixed_paths.py",
"InnerEye/Common/common_util.py",
]
for r in checkpoints_relative:
expected_files.append(f"{CHECKPOINT_FOLDER}/{r}")
for expected_file in expected_files:
assert (model_folder /
expected_file).is_file(), f"File missing: {expected_file}"
trm = model_folder / "TestsOutsidePackage/test_register_model.py"
if extra_code_directory:
assert trm.is_file()
else:
assert not trm.is_file()
def test_overridable_properties() -> None:
"""
Test to make sure all valid types can be parsed by the config parser
"""
overridable = [
"--num_dataload_workers=100", "--local_dataset=hello_world",
"--norm_method=Simple Norm", "--l_rate=100.0",
"--test_crop_size=1,2,3", "--output_range=-100.0,100.0"
]
parser = SegmentationModelBase.create_argparser()
args = vars(parser.parse_args(overridable))
assert args["num_dataload_workers"] == 100
assert str(args["local_dataset"]) == "hello_world"
assert args["norm_method"] == PhotometricNormalizationMethod.SimpleNorm
assert args["test_crop_size"] == (1, 2, 3)
assert args["output_range"] == (-100.0, 100.0)
def test_scale_and_unscale_image(
test_output_dirs: TestOutputDirectories) -> None:
"""
Test if an image in the CT value range can be recovered when we save dataset examples
(undoing the effects of CT Windowing)
"""
image_size = (5, 5, 5)
spacing = (1, 2, 3)
header = ImageHeader(origin=(0, 1, 0),
direction=(-1, 0, 0, 0, -1, 0, 0, 0, -1),
spacing=spacing)
np.random.seed(0)
# Random image values with mean -100, std 100. This will cover a range
# from -400 to +200 HU
image = np.random.normal(-100, 100, size=image_size)
window = 200
level = -100
# Lower and upper bounds of the interval of raw CT values that will be retained.
lower = level - window / 2
upper = level + window / 2
# Create a copy of the image with all values outside of the (Window, Level) range set to the boundaries.
# When saving and loading back in, we will not be able to recover any values that fell outside those boundaries.
image_restricted = image.copy()
image_restricted[image < lower] = lower
image_restricted[image > upper] = upper
# The image will be saved with voxel type short
image_restricted = image_restricted.astype(int)
# Apply window and level, mapping to the usual CNN input value range
cnn_input_range = (-1, +1)
image_windowed = LinearTransform.transform(data=image,
input_range=(lower, upper),
output_range=cnn_input_range)
args = SegmentationModelBase(
norm_method=PhotometricNormalizationMethod.CtWindow,
output_range=cnn_input_range,
window=window,
level=level,
should_validate=False)
file_name = test_output_dirs.create_file_or_folder_path(
"scale_and_unscale_image.nii.gz")
io_util.store_image_as_short_nifti(image_windowed, header, file_name, args)
image_from_disk = io_util.load_nifti_image(file_name)
# noinspection PyTypeChecker
assert_nifti_content(file_name, image_size, header,
np.unique(image_restricted).tolist(), np.short)
assert np.array_equal(image_from_disk.image, image_restricted)
def test_visualize_patch_sampling_2d(
test_output_dirs: TestOutputDirectories) -> None:
"""
Tests if patch sampling works for 2D images.
:param test_output_dirs:
"""
set_random_seed(0)
shape = (1, 20, 30)
foreground_classes = ["fg"]
class_weights = equally_weighted_classes(foreground_classes)
config = SegmentationModelBase(should_validate=False,
crop_size=(1, 5, 10),
class_weights=class_weights)
image = np.random.rand(1, *shape).astype(np.float32) * 1000
mask = np.ones(shape)
labels = np.zeros((len(class_weights), ) + shape)
labels[1, 0, 8:12, 5:25] = 1
labels[0] = 1 - labels[1]
output_folder = Path(test_output_dirs.root_dir)
image_header = None
sample = Sample(image=image,
mask=mask,
labels=labels,
metadata=PatientMetadata(patient_id='123',
image_header=image_header))
heatmap = visualize_random_crops(sample,
config,
output_folder=output_folder)
expected_folder = full_ml_test_data_path("patch_sampling")
expected_heatmap = expected_folder / "sampling_2d.npy"
# To update the stored results, uncomment this line:
# np.save(str(expected_heatmap), heatmap)
assert np.allclose(heatmap, np.load(
str(expected_heatmap))), "Patch sampling created a different heatmap."
assert len(list(output_folder.rglob("*.nii.gz"))) == 0
assert len(list(output_folder.rglob("*.png"))) == 1
actual_file = output_folder / "123_sampled_patches.png"
assert_file_exists(actual_file)
expected = expected_folder / "sampling_2d.png"
# To update the stored results, uncomment this line:
# expected.write_bytes(actual_file.read_bytes())
if not is_running_on_azure():
# When running on the Azure build agents, it appears that the bounding box of the images
# is slightly different than on local runs, even with equal dpi settings.
# It says: Image sizes don't match: actual (685, 469), expected (618, 424)
# Not able to figure out how to make the run results consistent, hence disable in cloud runs.
assert_binary_files_match(actual_file, expected)
def visualize_random_crops_for_dataset(
config: SegmentationModelBase,
output_folder: Optional[Path] = None) -> None:
"""
For segmentation models only: This function generates visualizations of the effect of sampling random patches
for training. Visualizations are stored in both Nifti format, and as 3 PNG thumbnail files, in the output folder.
:param config: The model configuration.
:param output_folder: The folder in which the visualizations should be written. If not provided, use a subfolder
"patch_sampling" in the models's default output folder
"""
dataset_splits = config.get_dataset_splits()
# Load a sample using the full image data loader
full_image_dataset = FullImageDataset(config, dataset_splits.train)
output_folder = output_folder or config.outputs_folder / PATCH_SAMPLING_FOLDER
count = min(config.show_patch_sampling, len(full_image_dataset))
for sample_index in range(count):
sample = full_image_dataset.get_samples_at_index(index=sample_index)[0]
visualize_random_crops(sample, config, output_folder=output_folder)
def test_invalid_stride_size() -> None:
config = SegmentationModelBase(
architecture="UNet3D",
feature_channels=[1],
crop_size=(64, 64, 64),
test_crop_size=(80, 80, 80),
image_channels=["mr"],
ground_truth_ids=["tumour_mass", "subtract"],
train_batch_size=8,
inference_batch_size=1,
inference_stride_size=(120, 120, 120),
should_validate=False
)
with pytest.raises(ValueError) as ex:
model_and_info = ModelAndInfo(config=config, model_execution_mode=ModelExecutionMode.TEST,
checkpoint_path=None)
model_and_info.try_create_model_load_from_checkpoint_and_adjust()
assert "inference stride size must be smaller" in ex.value.args[0]
assert str(config.inference_stride_size) in ex.value.args[0]
assert str(config.test_crop_size) in ex.value.args[0]
