From fade6c306e5d2ddc070df3fd135d8049028270c5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fernando=20P=C3=A9rez-Garc=C3=ADa?= Date: Thu, 3 Mar 2022 14:20:17 +0000 Subject: [PATCH 1/5] Ensure patches shape is compatible with model constraints Update CHANGELOG Fix transform being passed to itself Add correct changes to inference file --- CHANGELOG.md | 2 ++ InnerEye/ML/pipelines/inference.py | 43 ++++++++++++++++++++++++------ 2 files changed, 37 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6c0d3520f..16981cbe1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -90,6 +90,8 @@ gets uploaded to AzureML, by skipping all test folders. - ([#632](https://github.com/microsoft/InnerEye-DeepLearning/pull/632)) Nifti test data is no longer stored in Git LFS ### Fixed + +- ([#682](https://github.com/microsoft/InnerEye-DeepLearning/pull/682)) Ensure the shape of input patches is compatible with model constraints. - ([#659](https://github.com/microsoft/InnerEye-DeepLearning/pull/659)) Fix caching and checkpointing for TCGA CRCk dataset. - ([#649](https://github.com/microsoft/InnerEye-DeepLearning/pull/649)) Fix for the _convert_to_tensor_if_necessary method so that PIL.Image as well as np.array get converted to torch.Tensor. - ([#606](https://github.com/microsoft/InnerEye-DeepLearning/pull/606)) Bug fix: registered models do not include the hi-ml submodule diff --git a/InnerEye/ML/pipelines/inference.py b/InnerEye/ML/pipelines/inference.py index bfe3755f2..7b58680b0 100644 --- a/InnerEye/ML/pipelines/inference.py +++ b/InnerEye/ML/pipelines/inference.py @@ -7,7 +7,7 @@ import logging from enum import Enum from pathlib import Path -from typing import Optional, Tuple +from typing import Optional, Tuple, Dict import numpy as np import torch @@ -235,7 +235,7 @@ def post_process_posteriors(self, posteriors: np.ndarray, mask: np.ndarray = Non @torch.no_grad() def predict_whole_image(self, image_channels: np.ndarray, voxel_spacing_mm: TupleFloat3, - mask: np.ndarray = None, + mask: Optional[np.ndarray] = None, patient_id: int = 0) -> InferencePipeline.Result: """ Performs a single inference pass through the pipeline for the provided image @@ -255,12 +255,24 @@ def predict_whole_image(self, image_channels: np.ndarray, self.model.eval() image = tio.ScalarImage(tensor=image_channels) - subject = tio.Subject(image=image) + INPUT = 'input_image' + MASK = 'mask' + + subject_dict: Dict[str, tio.Image] = {INPUT: image} + if mask is not None: + subject_dict[MASK] = tio.LabelMap(tensor=mask[np.newaxis]) + subject = tio.Subject(subject_dict) + + constraints = self.model.model.crop_size_constraints + + # Make sure the image size is compatible with the model + ensure_shape_multiple = tio.EnsureShapeMultiple(constraints.multiple_of) + subject = ensure_shape_multiple(subject) # type: ignore # There may be cases where the test image is smaller than the test_crop_size. Adjust crop_size # to always fit into image. If test_crop_size is smaller than the image, crop will remain unchanged. - restrict_patch_size = self.model.model.crop_size_constraints.restrict_crop_size_to_image # type: ignore - effective_patch_size, effective_stride = restrict_patch_size(image.spatial_shape, # type: ignore + restrict_patch_size = constraints.restrict_crop_size_to_image # type: ignore + effective_patch_size, effective_stride = restrict_patch_size(subject.spatial_shape, # type: ignore self.model_config.test_crop_size, self.model_config.inference_stride_size) @@ -276,10 +288,10 @@ def predict_whole_image(self, image_channels: np.ndarray, aggregator = tio.inference.GridAggregator(grid_sampler) logging.debug( - f"Inference on image size {image.spatial_shape} will run " + f"Inference on image size {subject.spatial_shape} will run " f"with crop size {effective_patch_size} and stride {effective_stride}") for patches_batch in patch_loader: - input_tensor = patches_batch['image'][tio.DATA].float() + input_tensor = patches_batch[INPUT][tio.DATA].float() if self.model_config.use_gpu: input_tensor = input_tensor.cuda() locations = patches_batch[tio.LOCATION] @@ -288,9 +300,24 @@ def predict_whole_image(self, image_channels: np.ndarray, # collect the predictions over each of the batches aggregator.add_batch(patches_posteriors, locations) posteriors = aggregator.get_output_tensor().numpy() - posteriors, segmentation = self.post_process_posteriors(posteriors, mask=mask) + posteriors_mask = None if mask is None else subject[MASK].numpy()[0] + posteriors, segmentation = self.post_process_posteriors(posteriors, mask=posteriors_mask) image_util.check_array_range(posteriors, error_prefix="Whole image posteriors") + + # Make sure the final shape matches the input shape by undoing the padding in EnsureShapeMultiple (if any) + posteriors_image = tio.ScalarImage(tensor=posteriors, affine=image.affine) + segmentation_image = tio.LabelMap(tensor=segmentation[np.newaxis], affine=image.affine) + subject.add_image(posteriors_image, 'posteriors') + subject.add_image(segmentation_image, 'segmentation') + # Remove some images to avoid unnecessary computations + subject.remove_image(INPUT) + if mask is not None: + subject.remove_image(MASK) + subject_original_space = subject.apply_inverse_transform() + posteriors = subject_original_space.posteriors.numpy() # type: ignore + segmentation = subject_original_space.segmentation.numpy()[0] # type: ignore + # prepare pipeline results from the processed batch return InferencePipeline.Result( patient_id=patient_id, From e1354ac46e29d2ca0b93cf504d91fbe40d53e08d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fernando=20P=C3=A9rez-Garc=C3=ADa?= Date: Thu, 3 Mar 2022 15:22:09 +0000 Subject: [PATCH 2/5] Ignore mypy error --- InnerEye/ML/pipelines/inference.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/InnerEye/ML/pipelines/inference.py b/InnerEye/ML/pipelines/inference.py index 7b58680b0..5da98b2b3 100644 --- a/InnerEye/ML/pipelines/inference.py +++ b/InnerEye/ML/pipelines/inference.py @@ -266,7 +266,7 @@ def predict_whole_image(self, image_channels: np.ndarray, constraints = self.model.model.crop_size_constraints # Make sure the image size is compatible with the model - ensure_shape_multiple = tio.EnsureShapeMultiple(constraints.multiple_of) + ensure_shape_multiple = tio.EnsureShapeMultiple(constraints.multiple_of) # type: ignore subject = ensure_shape_multiple(subject) # type: ignore # There may be cases where the test image is smaller than the test_crop_size. Adjust crop_size From d47538bc43a48cbe86aa2f3af0c869e436379596 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fernando=20P=C3=A9rez-Garc=C3=ADa?= Date: Thu, 3 Mar 2022 17:40:40 +0000 Subject: [PATCH 3/5] Use TorchIO transform only when constraints are not None --- InnerEye/ML/pipelines/inference.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/InnerEye/ML/pipelines/inference.py b/InnerEye/ML/pipelines/inference.py index 5da98b2b3..32ac4c029 100644 --- a/InnerEye/ML/pipelines/inference.py +++ b/InnerEye/ML/pipelines/inference.py @@ -266,8 +266,10 @@ def predict_whole_image(self, image_channels: np.ndarray, constraints = self.model.model.crop_size_constraints # Make sure the image size is compatible with the model - ensure_shape_multiple = tio.EnsureShapeMultiple(constraints.multiple_of) # type: ignore - subject = ensure_shape_multiple(subject) # type: ignore + multiple_constraints = constraints.multiple_of + if multiple_constraints is not None: + ensure_shape_multiple = tio.EnsureShapeMultiple(constraints.multiple_of) # type: ignore + subject = ensure_shape_multiple(subject) # type: ignore # There may be cases where the test image is smaller than the test_crop_size. Adjust crop_size # to always fit into image. If test_crop_size is smaller than the image, crop will remain unchanged. @@ -314,7 +316,7 @@ def predict_whole_image(self, image_channels: np.ndarray, subject.remove_image(INPUT) if mask is not None: subject.remove_image(MASK) - subject_original_space = subject.apply_inverse_transform() + subject_original_space = subject.apply_inverse_transform() if subject.applied_transforms else subject posteriors = subject_original_space.posteriors.numpy() # type: ignore segmentation = subject_original_space.segmentation.numpy()[0] # type: ignore From f4c663400fb7a639f46ac5233ea525a487c65c18 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fernando=20P=C3=A9rez-Garc=C3=ADa?= Date: Thu, 3 Mar 2022 17:52:56 +0000 Subject: [PATCH 4/5] Ignore yet another mypy error --- InnerEye/ML/pipelines/inference.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/InnerEye/ML/pipelines/inference.py b/InnerEye/ML/pipelines/inference.py index 32ac4c029..b5e731c39 100644 --- a/InnerEye/ML/pipelines/inference.py +++ b/InnerEye/ML/pipelines/inference.py @@ -266,7 +266,7 @@ def predict_whole_image(self, image_channels: np.ndarray, constraints = self.model.model.crop_size_constraints # Make sure the image size is compatible with the model - multiple_constraints = constraints.multiple_of + multiple_constraints = constraints.multiple_of # type: ignore if multiple_constraints is not None: ensure_shape_multiple = tio.EnsureShapeMultiple(constraints.multiple_of) # type: ignore subject = ensure_shape_multiple(subject) # type: ignore From 29e62e016b5ff2c4eb26a0b9e0b51fe6928fde73 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fernando=20P=C3=A9rez-Garc=C3=ADa?= Date: Mon, 7 Mar 2022 10:25:27 +0000 Subject: [PATCH 5/5] Remove test for images that are too small --- Tests/ML/pipelines/test_inference_smallimages.py | 9 --------- 1 file changed, 9 deletions(-) diff --git a/Tests/ML/pipelines/test_inference_smallimages.py b/Tests/ML/pipelines/test_inference_smallimages.py index db2a4fcd9..d16a354c8 100644 --- a/Tests/ML/pipelines/test_inference_smallimages.py +++ b/Tests/ML/pipelines/test_inference_smallimages.py @@ -59,15 +59,6 @@ def run_inference_on_unet(size: TupleInt3) -> None: image_util.check_array_range(p) -def test_inference_on_too_small_image() -> None: - """ - Running inference on a simplified Unet model when the input image is too small along an axis. - """ - with pytest.raises(ValueError) as ex: - run_inference_on_unet((5, 10, 64)) - assert "input image must have at least a size of (16, 16, 16)" in str(ex) - - @pytest.mark.parametrize("size", [(26, 20, 50), (16, 16, 16)]) def test_inference_on_small_image(size: TupleInt3) -> None: """