torch_em.transform.augmentation

  1import torch
  2import numpy as np
  3import kornia
  4from skimage.transform import resize
  5
  6from ..util import ensure_tensor
  7
  8
  9class RandomElasticDeformationStacked(kornia.augmentation.AugmentationBase3D):
 10    def __init__(self,
 11                 control_point_spacing=1,
 12                 sigma=(32.0, 32.0),
 13                 alpha=(4.0, 4.0),
 14                 interpolation=kornia.constants.Resample.BILINEAR,
 15                 p=0.5,
 16                 keepdim=False,
 17                 same_on_batch=True):
 18        super().__init__(p=p,  # keepdim=keepdim,
 19                         same_on_batch=same_on_batch)
 20        if isinstance(control_point_spacing, int):
 21            self.control_point_spacing = [control_point_spacing] * 2
 22        else:
 23            self.control_point_spacing = control_point_spacing
 24        assert len(self.control_point_spacing) == 2
 25        self.interpolation = interpolation
 26        self.flags = dict(
 27            interpolation=torch.tensor(self.interpolation.value),
 28            sigma=sigma,
 29            alpha=alpha
 30        )
 31
 32    # The same transformation applied to all samples in a batch
 33    def generate_parameters(self, batch_shape):
 34        assert len(batch_shape) == 5
 35        shape = batch_shape[3:]
 36        control_shape = tuple(
 37            sh // spacing for sh, spacing in zip(shape, self.control_point_spacing)
 38        )
 39        deformation_fields = [
 40            np.random.uniform(-1, 1, control_shape),
 41            np.random.uniform(-1, 1, control_shape)
 42        ]
 43        deformation_fields = [
 44            resize(df, shape, order=3)[None] for df in deformation_fields
 45        ]
 46        noise = np.concatenate(deformation_fields, axis=0)[None].astype("float32")
 47        noise = torch.from_numpy(noise)
 48        return {"noise": noise}
 49
 50    def __call__(self, input, params=None):
 51        assert len(input.shape) == 5
 52        if params is None:
 53            params = self.generate_parameters(input.shape)
 54            self._params = params
 55
 56        noise = params["noise"]
 57        mode = "bilinear" if (self.flags["interpolation"] == 1).all() else "nearest"
 58        noise_ch = noise.expand(input.shape[1], -1, -1, -1)
 59        input_transformed = []
 60        for i, x in enumerate(torch.unbind(input, dim=0)):
 61            x_transformed = kornia.geometry.transform.elastic_transform2d(
 62                            x, noise_ch, sigma=self.flags["sigma"],
 63                            alpha=self.flags["alpha"], mode=mode,
 64                            padding_mode="reflection"
 65                            )
 66            input_transformed.append(x_transformed)
 67        input_transformed = torch.stack(input_transformed)
 68        return input_transformed
 69
 70
 71class RandomElasticDeformation(kornia.augmentation.AugmentationBase2D):
 72    def __init__(self,
 73                 control_point_spacing=1,
 74                 sigma=(4.0, 4.0),
 75                 alpha=(32.0, 32.0),
 76                 resample=kornia.constants.Resample.BILINEAR,
 77                 p=0.5,
 78                 keepdim=False,
 79                 same_on_batch=False):
 80        super().__init__(p=p,  # keepdim=keepdim,
 81                         same_on_batch=same_on_batch)
 82        if isinstance(control_point_spacing, int):
 83            self.control_point_spacing = [control_point_spacing] * 2
 84        else:
 85            self.control_point_spacing = control_point_spacing
 86        assert len(self.control_point_spacing) == 2
 87        self.resample = resample
 88        self.flags = dict(
 89            resample=torch.tensor(self.resample.value),
 90            sigma=sigma,
 91            alpha=alpha
 92        )
 93
 94    # TODO do we need special treatment for batches, channels > 1?
 95    def generate_parameters(self, batch_shape):
 96        assert len(batch_shape) == 4, f"{len(batch_shape)}"
 97        shape = batch_shape[2:]
 98        control_shape = tuple(
 99            sh // spacing for sh, spacing in zip(shape, self.control_point_spacing)
100        )
101        deformation_fields = [
102            np.random.uniform(-1, 1, control_shape),
103            np.random.uniform(-1, 1, control_shape)
104        ]
105        deformation_fields = [
106            resize(df, shape, order=3)[None] for df in deformation_fields
107        ]
108        noise = np.concatenate(deformation_fields, axis=0)[None].astype("float32")
109        noise = torch.from_numpy(noise)
110        return {"noise": noise}
111
112    def __call__(self, input, params=None):
113        if params is None:
114            params = self.generate_parameters(input.shape)
115            self._params = params
116        noise = params["noise"]
117        mode = "bilinear" if (self.flags["resample"] == 1).all() else "nearest"
118        return kornia.geometry.transform.elastic_transform2d(
119            input, noise, sigma=self.flags["sigma"], alpha=self.flags["alpha"], mode=mode,
120            padding_mode="reflection"
121        )
122
123
124# TODO implement 'require_halo', and estimate the halo from the transformations
125# so that we can load a bigger block and cut it away
126class KorniaAugmentationPipeline(torch.nn.Module):
127    interpolatable_torch_types = [torch.float16, torch.float32, torch.float64]
128    interpolatable_numpy_types = [np.dtype("float32"), np.dtype("float64")]
129
130    def __init__(self, *kornia_augmentations, dtype=torch.float32):
131        super().__init__()
132        self.augmentations = torch.nn.ModuleList(kornia_augmentations)
133        self.dtype = dtype
134        self.halo = self.compute_halo()
135
136    # for now we only add a halo for the random rotation trafos and
137    # also don't compute the halo dynamically based on the input shape
138    def compute_halo(self):
139        halo = None
140        for aug in self.augmentations:
141            if isinstance(aug, kornia.augmentation.RandomRotation):
142                halo = [32, 32]
143            if isinstance(aug, kornia.augmentation.RandomRotation3D):
144                halo = [32, 32, 32]
145        return halo
146
147    def is_interpolatable(self, tensor):
148        if torch.is_tensor(tensor):
149            return tensor.dtype in self.interpolatable_torch_types
150        else:
151            return tensor.dtype in self.interpolatable_numpy_types
152
153    def transform_tensor(self, augmentation, tensor, interpolatable, params=None):
154        interpolating = "interpolation" in getattr(augmentation, "flags", [])
155        if interpolating:
156            resampler = kornia.constants.Resample.get("BILINEAR" if interpolatable else "NEAREST")
157            augmentation.flags["interpolation"] = torch.tensor(resampler.value)
158        transformed = augmentation(tensor, params)
159        return transformed, augmentation._params
160
161    def forward(self, *tensors):
162        interpolatable = [self.is_interpolatable(tensor) for tensor in tensors]
163        tensors = [ensure_tensor(tensor, self.dtype) for tensor in tensors]
164        for aug in self.augmentations:
165
166            t0, params = self.transform_tensor(aug, tensors[0], interpolatable[0])
167            transformed_tensors = [t0]
168            for tensor, interpolate in zip(tensors[1:], interpolatable[1:]):
169                tensor, _ = self.transform_tensor(aug, tensor, interpolate, params=params)
170                transformed_tensors.append(tensor)
171
172            tensors = transformed_tensors
173        return tensors
174
175    def halo(self, shape):
176        return self.halo
177
178
179# TODO elastic deformation
180# Try out:
181# - RandomPerspective
182AUGMENTATIONS = {
183    "RandomAffine": {"degrees": 90, "scale": (0.9, 1.1)},
184    "RandomAffine3D": {"degrees": (90, 90, 90), "scale": (0.0, 1.1)},
185    "RandomDepthicalFlip3D": {},
186    "RandomHorizontalFlip": {},
187    "RandomHorizontalFlip3D": {},
188    "RandomRotation": {"degrees": 90},
189    "RandomRotation3D": {"degrees": (90, 90, 90)},
190    "RandomVerticalFlip": {},
191    "RandomVerticalFlip3D": {},
192    "RandomElasticDeformation3D": {"alpha": [5, 5], "sigma": [30, 30]}
193}
194
195
196DEFAULT_2D_AUGMENTATIONS = [
197    "RandomHorizontalFlip",
198    "RandomVerticalFlip"
199]
200DEFAULT_3D_AUGMENTATIONS = [
201    "RandomHorizontalFlip3D",
202    "RandomVerticalFlip3D",
203    "RandomDepthicalFlip3D",
204]
205DEFAULT_ANISOTROPIC_AUGMENTATIONS = [
206    "RandomHorizontalFlip3D",
207    "RandomVerticalFlip3D",
208    "RandomDepthicalFlip3D",
209]
210
211
212def create_augmentation(trafo):
213    assert trafo in dir(kornia.augmentation) or trafo in globals().keys(), f"Transformation {trafo} not defined"
214    if trafo in dir(kornia.augmentation):
215        return getattr(kornia.augmentation, trafo)(**AUGMENTATIONS[trafo])
216
217    return globals()[trafo](**AUGMENTATIONS[trafo])
218
219
220def get_augmentations(ndim=2,
221                      transforms=None,
222                      dtype=torch.float32):
223    if transforms is None:
224        assert ndim in (2, 3, "anisotropic"), f"Expect ndim to be one of (2, 3, 'anisotropic'), got {ndim}"
225        if ndim == 2:
226            transforms = DEFAULT_2D_AUGMENTATIONS
227        elif ndim == 3:
228            transforms = DEFAULT_3D_AUGMENTATIONS
229        else:
230            transforms = DEFAULT_ANISOTROPIC_AUGMENTATIONS
231    transforms = [create_augmentation(trafo) for trafo in transforms]
232
233    assert all(isinstance(trafo, kornia.augmentation.base._AugmentationBase)
234               for trafo in transforms)
235    augmentations = KorniaAugmentationPipeline(
236        *transforms,
237        dtype=dtype
238    )
239    return augmentations