[WIP] 365 Inverse spatial transforms

monai/data/__init__.py

@@ -25,12 +25,14 @@
 from .grid_dataset import GridPatchDataset, PatchDataset
 from .image_dataset import ImageDataset
 from .image_reader import ImageReader, ITKReader, NibabelReader, NumpyReader, PILReader
+from .inverse_batch_transform import BatchInverseTransform
 from .iterable_dataset import IterableDataset
 from .nifti_saver import NiftiSaver
 from .nifti_writer import write_nifti
 from .png_saver import PNGSaver
 from .png_writer import write_png
 from .synthetic import create_test_image_2d, create_test_image_3d
+from .test_time_augmentation import TestTimeAugmentation
 from .thread_buffer import ThreadBuffer
 from .utils import (
     DistributedSampler,

monai/data/image_reader.py

@@ -17,12 +17,10 @@
 import numpy as np
 from torch.utils.data._utils.collate import np_str_obj_array_pattern
 
+import monai.data.utils
 from monai.config import DtypeLike, KeysCollection
-from monai.data.utils import correct_nifti_header_if_necessary
 from monai.utils import ensure_tuple, optional_import
 
-from .utils import is_supported_format
-
 if TYPE_CHECKING:
     import itk  # type: ignore
     import nibabel as nib
@@ -322,7 +320,7 @@ def verify_suffix(self, filename: Union[Sequence[str], str]) -> bool:
 
         """
         suffixes: Sequence[str] = ["nii", "nii.gz"]
-        return has_nib and is_supported_format(filename, suffixes)
+        return has_nib and monai.data.is_supported_format(filename, suffixes)
 
     def read(self, data: Union[Sequence[str], str], **kwargs):
         """
@@ -343,7 +341,7 @@ def read(self, data: Union[Sequence[str], str], **kwargs):
         kwargs_.update(kwargs)
         for name in filenames:
             img = nib.load(name, **kwargs_)
-            img = correct_nifti_header_if_necessary(img)
+            img = monai.data.utils.correct_nifti_header_if_necessary(img)
             img_.append(img)
         return img_ if len(filenames) > 1 else img_[0]
 
@@ -453,7 +451,7 @@ def verify_suffix(self, filename: Union[Sequence[str], str]) -> bool:
                 if a list of files, verify all the suffixes.
         """
         suffixes: Sequence[str] = ["npz", "npy"]
-        return is_supported_format(filename, suffixes)
+        return monai.data.is_supported_format(filename, suffixes)
 
     def read(self, data: Union[Sequence[str], str], **kwargs):
         """
@@ -537,7 +535,7 @@ def verify_suffix(self, filename: Union[Sequence[str], str]) -> bool:
                 if a list of files, verify all the suffixes.
         """
         suffixes: Sequence[str] = ["png", "jpg", "jpeg", "bmp"]
-        return has_pil and is_supported_format(filename, suffixes)
+        return has_pil and monai.data.is_supported_format(filename, suffixes)
 
     def read(self, data: Union[Sequence[str], str, np.ndarray], **kwargs):
         """

monai/data/inverse_batch_transform.py

@@ -0,0 +1,90 @@
+# Copyright 2020 - 2021 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Callable, Dict, Hashable, Optional, Tuple
+
+import numpy as np
+from torch.utils.data.dataloader import DataLoader as TorchDataLoader
+
+from monai.data.dataloader import DataLoader
+from monai.data.dataset import Dataset
+from monai.data.utils import decollate_batch, pad_list_data_collate
+from monai.transforms.croppad.array import CenterSpatialCrop
+from monai.transforms.inverse_transform import InvertibleTransform
+from monai.utils import first
+from monai.utils.misc import ensure_tuple
+
+__all__ = ["BatchInverseTransform"]
+
+
+class _BatchInverseDataset(Dataset):
+    def __init__(
+        self,
+        data: Dict[str, Any],
+        transform: InvertibleTransform,
+        keys: Optional[Tuple[Hashable, ...]],
+        pad_collation_used: bool,
+    ) -> None:
+        self.data = decollate_batch(data)
+        self.invertible_transform = transform
+        self.keys = ensure_tuple(keys) if keys else None
+        self.pad_collation_used = pad_collation_used
+
+    def __getitem__(self, index: int) -> Dict[Hashable, np.ndarray]:
+        data = dict(self.data[index])
+        # If pad collation was used, then we need to undo this first
+        if self.pad_collation_used:
+            keys = self.keys or [key for key in data.keys() if str(key) + "_transforms" in data.keys()]
+            for key in keys:
+                transform_key = str(key) + "_transforms"
+                transform = data[transform_key][-1]
+                if transform["class"] == "SpatialPadd":
+                    data[key] = CenterSpatialCrop(transform["orig_size"])(data[key])
+                    # remove transform
+                    data[transform_key].pop()
+
+        return self.invertible_transform.inverse(data, self.keys)
+
+
+class BatchInverseTransform:
+    """something"""
+
+    def __init__(
+        self, transform: InvertibleTransform, loader: TorchDataLoader, collate_fn: Optional[Callable] = None
+    ) -> None:
+        """
+        Args:
+            transform: a callable data transform on input data.
+            loader: data loader used to generate the batch of data.
+            collate_fn: how to collate data after inverse transformations. Default will use the DataLoader's default
+                collation method. If returning images of different sizes, this will likely create an error (since the
+                collation will concatenate arrays, requiring them to be the same size). In this case, using
+                `collate_fn=lambda x: x` might solve the problem.
+        """
+        self.transform = transform
+        self.batch_size = loader.batch_size
+        self.num_workers = loader.num_workers
+        self.collate_fn = collate_fn
+        self.pad_collation_used = loader.collate_fn == pad_list_data_collate
+
+    def __call__(self, data: Dict[str, Any], keys: Optional[Tuple[Hashable, ...]] = None) -> Any:
+
+        inv_ds = _BatchInverseDataset(data, self.transform, keys, self.pad_collation_used)
+        inv_loader = DataLoader(
+            inv_ds, batch_size=self.batch_size, num_workers=self.num_workers, collate_fn=self.collate_fn
+        )
+        try:
+            return first(inv_loader)
+        except RuntimeError as re:
+            re_str = str(re)
+            if "stack expects each tensor to be equal size" in re_str:
+                re_str += "\nMONAI hint: try creating `BatchInverseTransform` with `collate_fn=lambda x: x`."
+            raise RuntimeError(re_str)

monai/data/nifti_saver.py

@@ -14,9 +14,9 @@
 import numpy as np
 import torch
 
+import monai.data.utils
 from monai.config import DtypeLike
 from monai.data.nifti_writer import write_nifti
-from monai.data.utils import create_file_basename
 from monai.utils import GridSampleMode, GridSamplePadMode
 from monai.utils import ImageMetaKey as Key
 
@@ -104,7 +104,7 @@ def save(self, data: Union[torch.Tensor, np.ndarray], meta_data: Optional[Dict]
         if isinstance(data, torch.Tensor):
             data = data.detach().cpu().numpy()
 
-        filename = create_file_basename(self.output_postfix, filename, self.output_dir)
+        filename = monai.data.utils.create_file_basename(self.output_postfix, filename, self.output_dir)
         filename = f"{filename}{self.output_ext}"
         # change data shape to be (channel, h, w, d)
         while len(data.shape) < 4:

monai/data/nifti_writer.py

@@ -14,8 +14,8 @@
 import numpy as np
 import torch
 
+import monai.data.utils
 from monai.config import DtypeLike
-from monai.data.utils import compute_shape_offset, to_affine_nd
 from monai.networks.layers import AffineTransform
 from monai.utils import GridSampleMode, GridSamplePadMode, optional_import
 
@@ -95,15 +95,15 @@ def write_nifti(
     sr = min(data.ndim, 3)
     if affine is None:
         affine = np.eye(4, dtype=np.float64)
-    affine = to_affine_nd(sr, affine)
+    affine = monai.data.utils.to_affine_nd(sr, affine)
 
     if target_affine is None:
         target_affine = affine
-    target_affine = to_affine_nd(sr, target_affine)
+    target_affine = monai.data.utils.to_affine_nd(sr, target_affine)
 
     if np.allclose(affine, target_affine, atol=1e-3):
         # no affine changes, save (data, affine)
-        results_img = nib.Nifti1Image(data.astype(output_dtype), to_affine_nd(3, target_affine))
+        results_img = nib.Nifti1Image(data.astype(output_dtype), monai.data.utils.to_affine_nd(3, target_affine))
         nib.save(results_img, file_name)
         return
 
@@ -115,7 +115,7 @@ def write_nifti(
     data = nib.orientations.apply_orientation(data, ornt_transform)
     _affine = affine @ nib.orientations.inv_ornt_aff(ornt_transform, data_shape)
     if np.allclose(_affine, target_affine, atol=1e-3) or not resample:
-        results_img = nib.Nifti1Image(data.astype(output_dtype), to_affine_nd(3, target_affine))
+        results_img = nib.Nifti1Image(data.astype(output_dtype), monai.data.utils.to_affine_nd(3, target_affine))
         nib.save(results_img, file_name)
         return
 
@@ -125,7 +125,7 @@ def write_nifti(
     )
     transform = np.linalg.inv(_affine) @ target_affine
     if output_spatial_shape is None:
-        output_spatial_shape, _ = compute_shape_offset(data.shape, _affine, target_affine)
+        output_spatial_shape, _ = monai.data.utils.compute_shape_offset(data.shape, _affine, target_affine)
     output_spatial_shape_ = list(output_spatial_shape) if output_spatial_shape is not None else []
     if data.ndim > 3:  # multi channel, resampling each channel
         while len(output_spatial_shape_) < 3:
@@ -151,6 +151,6 @@ def write_nifti(
         )
         data_np = data_torch.squeeze(0).squeeze(0).detach().cpu().numpy()
 
-    results_img = nib.Nifti1Image(data_np.astype(output_dtype), to_affine_nd(3, target_affine))
+    results_img = nib.Nifti1Image(data_np.astype(output_dtype), monai.data.utils.to_affine_nd(3, target_affine))
     nib.save(results_img, file_name)
     return

monai/data/png_saver.py

@@ -14,8 +14,8 @@
 import numpy as np
 import torch
 
+import monai.data.utils
 from monai.data.png_writer import write_png
-from monai.data.utils import create_file_basename
 from monai.utils import ImageMetaKey as Key
 from monai.utils import InterpolateMode
 
@@ -90,7 +90,7 @@ def save(self, data: Union[torch.Tensor, np.ndarray], meta_data: Optional[Dict]
         if isinstance(data, torch.Tensor):
             data = data.detach().cpu().numpy()
 
-        filename = create_file_basename(self.output_postfix, filename, self.output_dir)
+        filename = monai.data.utils.create_file_basename(self.output_postfix, filename, self.output_dir)
         filename = f"{filename}{self.output_ext}"
 
         if data.shape[0] == 1:

monai/data/test_time_augmentation.py

@@ -0,0 +1,116 @@
+# Copyright 2020 - 2021 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Dict
+
+import numpy as np
+import torch
+
+from monai.data.dataloader import DataLoader
+from monai.data.dataset import Dataset
+from monai.data.inverse_batch_transform import BatchInverseTransform
+from monai.data.utils import pad_list_data_collate
+from monai.transforms.compose import Compose
+from monai.transforms.inverse_transform import InvertibleTransform
+from monai.transforms.transform import Randomizable
+
+__all__ = ["TestTimeAugmentation"]
+
+
+def is_transform_rand(transform):
+    if not isinstance(transform, Compose):
+        return isinstance(transform, Randomizable)
+    # call recursively for each sub-transform
+    return any(is_transform_rand(t) for t in transform.transforms)
+
+
+class TestTimeAugmentation:
+    def __init__(
+        self,
+        transform: InvertibleTransform,
+        batch_size,
+        num_workers,
+        inferrer_fn,
+        device,
+    ) -> None:
+        self.transform = transform
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.inferrer_fn = inferrer_fn
+        self.device = device
+
+        # check that the transform has at least one random component
+        if not is_transform_rand(self.transform):
+            raise RuntimeError(
+                type(self).__name__
+                + " requires a `Randomizable` transform or a"
+                + " `Compose` containing at least one `Randomizable` transform."
+            )
+
+    def __call__(
+        self, data: Dict[str, Any], num_examples=10, image_key="image", label_key="label", return_full_data=False
+    ):
+        d = dict(data)
+
+        # check num examples is multiple of batch size
+        if num_examples % self.batch_size != 0:
+            raise ValueError("num_examples should be multiple of batch size.")
+
+        # generate batch of data of size == batch_size, dataset and dataloader
+        data_in = [d for _ in range(num_examples)]
+        ds = Dataset(data_in, self.transform)
+        dl = DataLoader(ds, self.num_workers, batch_size=self.batch_size, collate_fn=pad_list_data_collate)
+
+        label_transform_key = label_key + "_transforms"
+
+        # create inverter
+        inverter = BatchInverseTransform(self.transform, dl)
+
+        outputs = []
+
+        for batch_data in dl:
+
+            batch_images = batch_data[image_key].to(self.device)
+
+            # do model forward pass
+            batch_output = self.inferrer_fn(batch_images)
+            if isinstance(batch_output, torch.Tensor):
+                batch_output = batch_output.detach().cpu()
+            if isinstance(batch_output, np.ndarray):
+                batch_output = torch.Tensor(batch_output)
+
+            # check binary labels are extracted
+            if not all(torch.unique(batch_output.int()) == torch.Tensor([0, 1])):
+                raise RuntimeError(
+                    "Test-time augmentation requires binary channels. If this is "
+                    "not binary segmentation, then you should one-hot your output."
+                )
+
+            # create a dictionary containing the inferred batch and their transforms
+            inferred_dict = {label_key: batch_output, label_transform_key: batch_data[label_transform_key]}
+
+            # do inverse transformation (only for the label key)
+            inv_batch = inverter(inferred_dict, label_key)
+
+            # append
+            outputs.append(inv_batch)
+
+        # calculate mean and standard deviation
+        output = np.concatenate(outputs)
+
+        if return_full_data:
+            return output
+
+        mode = np.apply_along_axis(lambda x: np.bincount(x).argmax(), axis=0, arr=output.astype(np.int64))
+        mean = np.mean(output, axis=0)
+        std = np.std(output, axis=0)
+        vvc = np.std(output) / np.mean(output)
+        return mode, mean, std, vvc