Pad model outputs if they are smaller than the inputs

CHANGELOG.md

@@ -90,6 +90,7 @@ 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
+- ([#681](https://github.com/microsoft/InnerEye-DeepLearning/pull/681)) Pad model outputs if they are smaller than the inputs.
 - ([#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

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()