diff --git a/CHANGELOG.md b/CHANGELOG.md index c4250ccd6..099962c97 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -93,6 +93,7 @@ gets uploaded to AzureML, by skipping all test folders. ### Fixed +- ([#681](https://github.com/microsoft/InnerEye-DeepLearning/pull/681)) Pad model outputs if they are smaller than the inputs. - ([#683](https://github.com/microsoft/InnerEye-DeepLearning/pull/683)) Fix missing separator error in docs Makefile. - ([#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. diff --git a/InnerEye/ML/pipelines/inference.py b/InnerEye/ML/pipelines/inference.py index bfe3755f2..5b397c6e0 100644 --- a/InnerEye/ML/pipelines/inference.py +++ b/InnerEye/ML/pipelines/inference.py @@ -285,6 +285,14 @@ def predict_whole_image(self, image_channels: np.ndarray, locations = patches_batch[tio.LOCATION] # perform the forward pass patches_posteriors = self.model(input_tensor).detach() + # pad posteriors if they are smaller than the input + input_shape = input_tensor.shape[-3:] + patches_posteriors_shape = patches_posteriors.shape[-3:] + if input_shape != patches_posteriors_shape: + difference = np.array(input_shape) - np.array(patches_posteriors_shape) + assert not np.any(difference % 2) # the differences in shape are expected to be even + padding = tuple(np.repeat(difference // 2, 2)) + patches_posteriors = torch.nn.functional.pad(patches_posteriors, padding) # collect the predictions over each of the batches aggregator.add_batch(patches_posteriors, locations) posteriors = aggregator.get_output_tensor().numpy()