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_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_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.adjust_after_mixed_precision_and_parallel(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_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.adjust_after_mixed_precision_and_parallel(model)
assert model_config.inference_stride_size == test_stride_size
model_config.inference_stride_size = None
model_config.adjust_after_mixed_precision_and_parallel(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_inference_identity(image_size: Any,
crop_size: Any,
shrink_by: Any,
num_classes: int,
create_mask: bool,
extract_largest_foreground_connected_component: bool,
is_ensemble: bool,
posterior_smoothing_mm: Any) -> None:
"""
Test to make sure inference pipeline is identity preserving, ie: we can recreate deterministic
model output, ensuring the patching and stitching is robust.
"""
# fix random seed
np.random.seed(0)
ground_truth_ids = list(map(str, range(num_classes)))
# image to run inference on: The mock model passes the input image through, hence the input
# image must have as many channels as we have classes (plus background), such that the output is
# also a valid posterior.
num_channels = num_classes + 1
image_channels = np.random.randn(num_channels, *list(image_size))
# create a random mask if required
mask = np.round(np.random.uniform(size=image_size)).astype(np.int) if create_mask else None
model_config = SegmentationModelBase(
crop_size=crop_size,
image_channels=list(map(str, range(num_channels))),
ground_truth_ids=ground_truth_ids,
should_validate=False,
posterior_smoothing_mm=posterior_smoothing_mm
)
# We have to set largest_connected_component_foreground_classes after creating the model config,
# because this parameter is not overridable and hence will not be set by GenericConfig's constructor.
if extract_largest_foreground_connected_component:
model_config.largest_connected_component_foreground_classes = [(c, None) for c in ground_truth_ids]
# set expected posteriors
expected_posteriors = torch.nn.functional.softmax(torch.tensor(image_channels), dim=0).numpy()
# apply the mask if required
if mask is not None:
expected_posteriors = image_util.apply_mask_to_posteriors(expected_posteriors, mask)
if posterior_smoothing_mm is not None:
expected_posteriors = image_util.gaussian_smooth_posteriors(
posteriors=expected_posteriors,
kernel_size_mm=posterior_smoothing_mm,
voxel_spacing_mm=(1, 1, 1)
)
# compute expected segmentation
expected_segmentation = image_util.posteriors_to_segmentation(expected_posteriors)
if extract_largest_foreground_connected_component:
largest_component = image_util.extract_largest_foreground_connected_component(
multi_label_array=expected_segmentation)
# make sure the test data is accurate by checking if more than one component exists
assert not np.array_equal(largest_component, expected_segmentation)
expected_segmentation = largest_component
# instantiate the model
model = PyTorchMockModel(shrink_by)
model_config.adjust_after_mixed_precision_and_parallel(model)
# create single or ensemble inference pipeline
inference_pipeline = InferencePipeline(model=model, model_config=model_config)
full_image_inference_pipeline = EnsemblePipeline([inference_pipeline], model_config) \
if is_ensemble else inference_pipeline
# compute full image inference results
inference_result = full_image_inference_pipeline \
.predict_and_post_process_whole_image(image_channels=image_channels, mask=mask, voxel_spacing_mm=(1, 1, 1))
# Segmentation must have same size as input image
assert inference_result.segmentation.shape == image_size
assert inference_result.posteriors.shape == (num_classes + 1,) + image_size
# check that the posteriors and segmentations are as expected. Flatten to a list so that the error
# messages are more informative.
assert np.allclose(inference_result.posteriors, expected_posteriors)
assert np.array_equal(inference_result.segmentation, expected_segmentation)
