torch_em.transform.invertible_augmentations

View Source

  1from typing import Callable
  2import torch
  3import kornia.augmentation as K
  4
  5
  6DEFAULT_WEAK_AUGMENTATIONS = {
  7    "intensity": {},
  8    "geometrical": {
  9        "RandomHorizontalFlip": {},
 10        "RandomVerticalFlip": {},
 11        "RandomRotation90": {"times": (-1, 2)},
 12    }
 13}
 14
 15DEFAULT_STRONG_AUGMENTATIONS = {
 16    "intensity": {
 17        "RandomGaussianBlur": {"kernel_size": (3, 3), "sigma": (0.1, 1.0)},
 18        "RandomGaussianNoise": {"mean": (0.0), "std": (0.1)},
 19    },
 20    "geometrical": {
 21        "RandomHorizontalFlip": {},
 22        "RandomVerticalFlip": {},
 23        "RandomRotation90": {"times": (-1, 2)},
 24    }
 25}
 26
 27
 28def get_intensity_augmentations(ndim, aug_name: str = None, aug_dict: dict = None, p: float = 0.75) -> callable:
 29    assert ndim == 2 or ndim == 3, f"Number of dimensions must be 2 or 3, got ndim={ndim}"
 30    assert aug_name is not None or aug_dict is not None, "Either aug_name or aug_dict must be provided."
 31
 32    if aug_dict is None:
 33        if aug_name == "weak":
 34            aug_dict = DEFAULT_WEAK_AUGMENTATIONS["intensity"]
 35        elif aug_name == "strong":
 36            aug_dict = DEFAULT_STRONG_AUGMENTATIONS["intensity"]
 37        else:
 38            raise ValueError(f"Augmentation name needs to be \"weak\" or \"strong\", got aug_name={aug_name}")
 39
 40    transforms_list = []
 41    for trafo, kwargs in aug_dict.items():
 42        assert trafo in dir(K), f"{trafo} not found in kornia.augmentation"
 43        transforms_list.append(getattr(K, trafo)(p=p, **kwargs))
 44
 45    if ndim == 2:
 46        return K.AugmentationSequential(*transforms_list, data_keys=["input"], same_on_batch=False)
 47    elif ndim == 3:
 48        return AugmentationSequential3D(*transforms_list)
 49
 50
 51def get_geometrical_augmentations(ndim, aug_name: str = None, aug_dict: dict = None, p: float = 0.75) -> callable:
 52    assert ndim == 2 or ndim == 3, f"Number of dimensions must be 2 or 3, got ndim={ndim}"
 53    assert aug_name is not None or aug_dict is not None, "Either aug_name or aug_dict must be provided."
 54
 55    if aug_dict is None:
 56        if aug_name == "weak":
 57            aug_dict = DEFAULT_WEAK_AUGMENTATIONS["geometrical"]
 58        elif aug_name == "strong":
 59            aug_dict = DEFAULT_STRONG_AUGMENTATIONS["geometrical"]
 60        else:
 61            raise ValueError(f"Augmentation name needs to be \"weak\" or \"strong\", got aug_name={aug_name}")
 62
 63    transforms_list = []
 64    for trafo, kwargs in aug_dict.items():
 65        assert trafo in dir(K), f"{trafo} not found in kornia.augmentation"
 66        transforms_list.append(getattr(K, trafo)(p=p, **kwargs))
 67
 68    if ndim == 2:
 69        return K.AugmentationSequential(*transforms_list, data_keys=["input"], same_on_batch=False)
 70    elif ndim == 3:
 71        return AugmentationSequential3D(*transforms_list)
 72
 73
 74def get_default_augmentations(aug_name: str, ndim: int, p: float = 0.75):
 75    assert ndim == 2 or ndim == 3, f"Number of dimensions must be 2 or 3, got ndim={ndim}"
 76
 77    if aug_name == "weak":
 78        intensity_transforms = get_intensity_augmentations(ndim, aug_name=aug_name, p=p)
 79        geometrical_transforms = get_geometrical_augmentations(ndim, aug_name=aug_name, p=p)
 80    elif aug_name == "strong":
 81        intensity_transforms = get_intensity_augmentations(ndim, aug_name=aug_name, p=p)
 82        geometrical_transforms = get_geometrical_augmentations(ndim, aug_name=aug_name, p=p)
 83    else:
 84        raise ValueError(f"Augmentation name needs to be \"weak\" or \"strong\", got aug_name={aug_name}")
 85
 86    return intensity_transforms, geometrical_transforms
 87
 88
 89class AugmentationSequential3D(torch.nn.Module):
 90    def __init__(self, *augmentations: torch.nn.Module):
 91        super().__init__()
 92        self.augmentations = torch.nn.ModuleList(augmentations)
 93        self._params = None
 94
 95    @staticmethod
 96    def _flatten(x):
 97        """
 98        (B, C, D, H, W) -> (B, C*D, H, W)
 99        """
100        if x.ndim != 5:
101            raise RuntimeError(f"Expected 5D tensor, got {x.shape}")
102        b, c, d, h, w = x.shape
103        x = x.reshape(b, c * d, h, w)
104        return x, (b, c, d, h, w)
105
106    @staticmethod
107    def _unflatten(x, shape):
108        """
109        (B, C*D, H, W) -> (B, C, D, H, W)
110        """
111        b, c, d, h, w = shape
112        x = x.reshape(b, c, d, h, w)
113        return x
114
115    def forward(self, x: torch.Tensor) -> torch.Tensor:
116        params_all = []
117
118        flat_x, shape = self._flatten(x)
119        for aug in self.augmentations:
120            flat_x = aug(flat_x)
121            params_all.append(aug._params)
122        out = self._unflatten(flat_x, shape)
123        self._params = params_all
124        return out
125
126    def inverse(self, x: torch.Tensor, params) -> torch.Tensor:
127
128        flat_x, shape = self._flatten(x)
129        for aug, p in reversed(list(zip(self.augmentations, params))):
130            flat_x = aug.inverse(flat_x, params=p)
131        out = self._unflatten(flat_x, shape)
132
133        return out
134
135
136class InvertibleAugmenter(torch.nn.Module):
137
138    def __init__(
139        self,
140        intensity_transforms: Callable[[torch.Tensor], torch.Tensor],
141        geometrical_transforms: Callable[[torch.Tensor], torch.Tensor],
142        clip_max = None,
143        **kwargs,
144    ):
145        super().__init__(**kwargs)
146        self.intensity_transforms = intensity_transforms
147        self.geometrical_transforms = geometrical_transforms
148        self.clip_max = clip_max
149
150    def reset(self):
151        self.params = None
152
153    def transform(self, x: torch.Tensor) -> torch.Tensor:
154        x = self.intensity_transforms(x)
155        x = self.geometrical_transforms(x)
156        # Clamp after the geometric transforms, because their interpolation can overshoot clip_max slightly.
157        if self.clip_max is not None:
158            x = torch.clamp(x, 0.0, self.clip_max)
159
160        self.params = self.geometrical_transforms._params
161
162        return x
163
164    def reverse_transform(self, x: torch.Tensor) -> torch.Tensor:
165        x_inv = self.geometrical_transforms.inverse(x, params=self.params)
166        return x_inv
167
168
169class MeanTeacherAugmenters:
170    def __init__(
171        self,
172        ndim: int,
173        teacher=None,
174        student=None,
175        clip_max=None,
176        aug_dict=None,
177    ):
178        self.teacher = teacher or self._get_weak_augmenter(ndim, clip_max, aug_dict)
179        self.student = student or self._get_weak_augmenter(ndim, clip_max, aug_dict)
180
181    def _get_weak_augmenter(self, ndim, clip_max, aug_dict):
182        if aug_dict is None:
183            augmenter = InvertibleAugmenter(
184                *get_default_augmentations("weak", ndim=ndim),
185                clip_max=clip_max,
186            )
187        else:
188            missing_keys = {"intensity", "geometrical"} - set(aug_dict)
189            if missing_keys:
190                raise ValueError(f"aug_dict needs the keys 'intensity' and 'geometrical', missing: {missing_keys}")
191
192            augmenter = InvertibleAugmenter(
193                get_intensity_augmentations(ndim, aug_dict=aug_dict["intensity"]),
194                get_geometrical_augmentations(ndim, aug_dict=aug_dict["geometrical"]),
195                clip_max=clip_max,
196            )
197        return augmenter
198
199    def reset_all(self):
200        self.teacher.reset()
201        self.student.reset()
202
203
204class FixMatchAugmenters:
205    def __init__(
206        self,
207        ndim: int,
208        teacher=None,
209        student=None,
210        clip_max=None,
211    ):
212        self.teacher = teacher or InvertibleAugmenter(*get_default_augmentations("weak", ndim=ndim), clip_max=clip_max)
213        self.student = student or InvertibleAugmenter(*get_default_augmentations("strong", ndim=ndim), clip_max=clip_max)
214
215    def reset_all(self):
216        self.teacher.reset()
217        self.student.reset()
218
219
220class UniMatchv2Augmenters:
221    def __init__(
222        self,
223        ndim: int,
224        weak=None,
225        strong1=None,
226        strong2=None,
227        clip_max=None,
228    ):
229        self.weak = weak or InvertibleAugmenter(*get_default_augmentations("weak", ndim=ndim), clip_max=clip_max)
230        self.strong1 = strong1 or InvertibleAugmenter(*get_default_augmentations("strong", ndim=ndim), clip_max=clip_max)
231        self.strong2 = strong2 or InvertibleAugmenter(*get_default_augmentations("strong", ndim=ndim), clip_max=clip_max)
232
233    def reset_all(self):
234        self.weak.reset()
235        self.strong1.reset()
236        self.strong2.reset()

DEFAULT_WEAK_AUGMENTATIONS = {'intensity': {}, 'geometrical': {'RandomHorizontalFlip': {}, 'RandomVerticalFlip': {}, 'RandomRotation90': {'times': (-1, 2)}}}

DEFAULT_STRONG_AUGMENTATIONS = {'intensity': {'RandomGaussianBlur': {'kernel_size': (3, 3), 'sigma': (0.1, 1.0)}, 'RandomGaussianNoise': {'mean': 0.0, 'std': 0.1}}, 'geometrical': {'RandomHorizontalFlip': {}, 'RandomVerticalFlip': {}, 'RandomRotation90': {'times': (-1, 2)}}}

def get_intensity_augmentations( ndim, aug_name: str = None, aug_dict: dict = None, p: float = 0.75) -> <built-in function callable>: View Source

29def get_intensity_augmentations(ndim, aug_name: str = None, aug_dict: dict = None, p: float = 0.75) -> callable:
30    assert ndim == 2 or ndim == 3, f"Number of dimensions must be 2 or 3, got ndim={ndim}"
31    assert aug_name is not None or aug_dict is not None, "Either aug_name or aug_dict must be provided."
32
33    if aug_dict is None:
34        if aug_name == "weak":
35            aug_dict = DEFAULT_WEAK_AUGMENTATIONS["intensity"]
36        elif aug_name == "strong":
37            aug_dict = DEFAULT_STRONG_AUGMENTATIONS["intensity"]
38        else:
39            raise ValueError(f"Augmentation name needs to be \"weak\" or \"strong\", got aug_name={aug_name}")
40
41    transforms_list = []
42    for trafo, kwargs in aug_dict.items():
43        assert trafo in dir(K), f"{trafo} not found in kornia.augmentation"
44        transforms_list.append(getattr(K, trafo)(p=p, **kwargs))
45
46    if ndim == 2:
47        return K.AugmentationSequential(*transforms_list, data_keys=["input"], same_on_batch=False)
48    elif ndim == 3:
49        return AugmentationSequential3D(*transforms_list)

def get_geometrical_augmentations( ndim, aug_name: str = None, aug_dict: dict = None, p: float = 0.75) -> <built-in function callable>: View Source

52def get_geometrical_augmentations(ndim, aug_name: str = None, aug_dict: dict = None, p: float = 0.75) -> callable:
53    assert ndim == 2 or ndim == 3, f"Number of dimensions must be 2 or 3, got ndim={ndim}"
54    assert aug_name is not None or aug_dict is not None, "Either aug_name or aug_dict must be provided."
55
56    if aug_dict is None:
57        if aug_name == "weak":
58            aug_dict = DEFAULT_WEAK_AUGMENTATIONS["geometrical"]
59        elif aug_name == "strong":
60            aug_dict = DEFAULT_STRONG_AUGMENTATIONS["geometrical"]
61        else:
62            raise ValueError(f"Augmentation name needs to be \"weak\" or \"strong\", got aug_name={aug_name}")
63
64    transforms_list = []
65    for trafo, kwargs in aug_dict.items():
66        assert trafo in dir(K), f"{trafo} not found in kornia.augmentation"
67        transforms_list.append(getattr(K, trafo)(p=p, **kwargs))
68
69    if ndim == 2:
70        return K.AugmentationSequential(*transforms_list, data_keys=["input"], same_on_batch=False)
71    elif ndim == 3:
72        return AugmentationSequential3D(*transforms_list)

def get_default_augmentations(aug_name: str, ndim: int, p: float = 0.75): View Source

75def get_default_augmentations(aug_name: str, ndim: int, p: float = 0.75):
76    assert ndim == 2 or ndim == 3, f"Number of dimensions must be 2 or 3, got ndim={ndim}"
77
78    if aug_name == "weak":
79        intensity_transforms = get_intensity_augmentations(ndim, aug_name=aug_name, p=p)
80        geometrical_transforms = get_geometrical_augmentations(ndim, aug_name=aug_name, p=p)
81    elif aug_name == "strong":
82        intensity_transforms = get_intensity_augmentations(ndim, aug_name=aug_name, p=p)
83        geometrical_transforms = get_geometrical_augmentations(ndim, aug_name=aug_name, p=p)
84    else:
85        raise ValueError(f"Augmentation name needs to be \"weak\" or \"strong\", got aug_name={aug_name}")
86
87    return intensity_transforms, geometrical_transforms

class AugmentationSequential3D(torch.nn.modules.module.Module): View Source

 90class AugmentationSequential3D(torch.nn.Module):
 91    def __init__(self, *augmentations: torch.nn.Module):
 92        super().__init__()
 93        self.augmentations = torch.nn.ModuleList(augmentations)
 94        self._params = None
 95
 96    @staticmethod
 97    def _flatten(x):
 98        """
 99        (B, C, D, H, W) -> (B, C*D, H, W)
100        """
101        if x.ndim != 5:
102            raise RuntimeError(f"Expected 5D tensor, got {x.shape}")
103        b, c, d, h, w = x.shape
104        x = x.reshape(b, c * d, h, w)
105        return x, (b, c, d, h, w)
106
107    @staticmethod
108    def _unflatten(x, shape):
109        """
110        (B, C*D, H, W) -> (B, C, D, H, W)
111        """
112        b, c, d, h, w = shape
113        x = x.reshape(b, c, d, h, w)
114        return x
115
116    def forward(self, x: torch.Tensor) -> torch.Tensor:
117        params_all = []
118
119        flat_x, shape = self._flatten(x)
120        for aug in self.augmentations:
121            flat_x = aug(flat_x)
122            params_all.append(aug._params)
123        out = self._unflatten(flat_x, shape)
124        self._params = params_all
125        return out
126
127    def inverse(self, x: torch.Tensor, params) -> torch.Tensor:
128
129        flat_x, shape = self._flatten(x)
130        for aug, p in reversed(list(zip(self.augmentations, params))):
131            flat_x = aug.inverse(flat_x, params=p)
132        out = self._unflatten(flat_x, shape)
133
134        return out

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call to(), etc.

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

AugmentationSequential3D(*augmentations: torch.nn.modules.module.Module) View Source

91    def __init__(self, *augmentations: torch.nn.Module):
92        super().__init__()
93        self.augmentations = torch.nn.ModuleList(augmentations)
94        self._params = None

Initialize internal Module state, shared by both nn.Module and ScriptModule.

augmentations

def forward(self, x: torch.Tensor) -> torch.Tensor: View Source

116    def forward(self, x: torch.Tensor) -> torch.Tensor:
117        params_all = []
118
119        flat_x, shape = self._flatten(x)
120        for aug in self.augmentations:
121            flat_x = aug(flat_x)
122            params_all.append(aug._params)
123        out = self._unflatten(flat_x, shape)
124        self._params = params_all
125        return out

Define the computation performed at every call.

Should be overridden by all subclasses.

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

def inverse(self, x: torch.Tensor, params) -> torch.Tensor: View Source

127    def inverse(self, x: torch.Tensor, params) -> torch.Tensor:
128
129        flat_x, shape = self._flatten(x)
130        for aug, p in reversed(list(zip(self.augmentations, params))):
131            flat_x = aug.inverse(flat_x, params=p)
132        out = self._unflatten(flat_x, shape)
133
134        return out

class InvertibleAugmenter(torch.nn.modules.module.Module): View Source

137class InvertibleAugmenter(torch.nn.Module):
138
139    def __init__(
140        self,
141        intensity_transforms: Callable[[torch.Tensor], torch.Tensor],
142        geometrical_transforms: Callable[[torch.Tensor], torch.Tensor],
143        clip_max = None,
144        **kwargs,
145    ):
146        super().__init__(**kwargs)
147        self.intensity_transforms = intensity_transforms
148        self.geometrical_transforms = geometrical_transforms
149        self.clip_max = clip_max
150
151    def reset(self):
152        self.params = None
153
154    def transform(self, x: torch.Tensor) -> torch.Tensor:
155        x = self.intensity_transforms(x)
156        x = self.geometrical_transforms(x)
157        # Clamp after the geometric transforms, because their interpolation can overshoot clip_max slightly.
158        if self.clip_max is not None:
159            x = torch.clamp(x, 0.0, self.clip_max)
160
161        self.params = self.geometrical_transforms._params
162
163        return x
164
165    def reverse_transform(self, x: torch.Tensor) -> torch.Tensor:
166        x_inv = self.geometrical_transforms.inverse(x, params=self.params)
167        return x_inv

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call to(), etc.

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

InvertibleAugmenter( intensity_transforms: Callable[[torch.Tensor], torch.Tensor], geometrical_transforms: Callable[[torch.Tensor], torch.Tensor], clip_max=None, **kwargs) View Source

139    def __init__(
140        self,
141        intensity_transforms: Callable[[torch.Tensor], torch.Tensor],
142        geometrical_transforms: Callable[[torch.Tensor], torch.Tensor],
143        clip_max = None,
144        **kwargs,
145    ):
146        super().__init__(**kwargs)
147        self.intensity_transforms = intensity_transforms
148        self.geometrical_transforms = geometrical_transforms
149        self.clip_max = clip_max

Initialize internal Module state, shared by both nn.Module and ScriptModule.

intensity_transforms

geometrical_transforms

clip_max

def reset(self): View Source

151    def reset(self):
152        self.params = None

def transform(self, x: torch.Tensor) -> torch.Tensor: View Source

154    def transform(self, x: torch.Tensor) -> torch.Tensor:
155        x = self.intensity_transforms(x)
156        x = self.geometrical_transforms(x)
157        # Clamp after the geometric transforms, because their interpolation can overshoot clip_max slightly.
158        if self.clip_max is not None:
159            x = torch.clamp(x, 0.0, self.clip_max)
160
161        self.params = self.geometrical_transforms._params
162
163        return x

def reverse_transform(self, x: torch.Tensor) -> torch.Tensor: View Source

165    def reverse_transform(self, x: torch.Tensor) -> torch.Tensor:
166        x_inv = self.geometrical_transforms.inverse(x, params=self.params)
167        return x_inv

class MeanTeacherAugmenters: View Source

170class MeanTeacherAugmenters:
171    def __init__(
172        self,
173        ndim: int,
174        teacher=None,
175        student=None,
176        clip_max=None,
177        aug_dict=None,
178    ):
179        self.teacher = teacher or self._get_weak_augmenter(ndim, clip_max, aug_dict)
180        self.student = student or self._get_weak_augmenter(ndim, clip_max, aug_dict)
181
182    def _get_weak_augmenter(self, ndim, clip_max, aug_dict):
183        if aug_dict is None:
184            augmenter = InvertibleAugmenter(
185                *get_default_augmentations("weak", ndim=ndim),
186                clip_max=clip_max,
187            )
188        else:
189            missing_keys = {"intensity", "geometrical"} - set(aug_dict)
190            if missing_keys:
191                raise ValueError(f"aug_dict needs the keys 'intensity' and 'geometrical', missing: {missing_keys}")
192
193            augmenter = InvertibleAugmenter(
194                get_intensity_augmentations(ndim, aug_dict=aug_dict["intensity"]),
195                get_geometrical_augmentations(ndim, aug_dict=aug_dict["geometrical"]),
196                clip_max=clip_max,
197            )
198        return augmenter
199
200    def reset_all(self):
201        self.teacher.reset()
202        self.student.reset()

MeanTeacherAugmenters(ndim: int, teacher=None, student=None, clip_max=None, aug_dict=None) View Source

171    def __init__(
172        self,
173        ndim: int,
174        teacher=None,
175        student=None,
176        clip_max=None,
177        aug_dict=None,
178    ):
179        self.teacher = teacher or self._get_weak_augmenter(ndim, clip_max, aug_dict)
180        self.student = student or self._get_weak_augmenter(ndim, clip_max, aug_dict)

teacher

student

def reset_all(self): View Source

200    def reset_all(self):
201        self.teacher.reset()
202        self.student.reset()

class FixMatchAugmenters: View Source

205class FixMatchAugmenters:
206    def __init__(
207        self,
208        ndim: int,
209        teacher=None,
210        student=None,
211        clip_max=None,
212    ):
213        self.teacher = teacher or InvertibleAugmenter(*get_default_augmentations("weak", ndim=ndim), clip_max=clip_max)
214        self.student = student or InvertibleAugmenter(*get_default_augmentations("strong", ndim=ndim), clip_max=clip_max)
215
216    def reset_all(self):
217        self.teacher.reset()
218        self.student.reset()

FixMatchAugmenters(ndim: int, teacher=None, student=None, clip_max=None) View Source

206    def __init__(
207        self,
208        ndim: int,
209        teacher=None,
210        student=None,
211        clip_max=None,
212    ):
213        self.teacher = teacher or InvertibleAugmenter(*get_default_augmentations("weak", ndim=ndim), clip_max=clip_max)
214        self.student = student or InvertibleAugmenter(*get_default_augmentations("strong", ndim=ndim), clip_max=clip_max)

teacher

student

def reset_all(self): View Source

216    def reset_all(self):
217        self.teacher.reset()
218        self.student.reset()

class UniMatchv2Augmenters: View Source

221class UniMatchv2Augmenters:
222    def __init__(
223        self,
224        ndim: int,
225        weak=None,
226        strong1=None,
227        strong2=None,
228        clip_max=None,
229    ):
230        self.weak = weak or InvertibleAugmenter(*get_default_augmentations("weak", ndim=ndim), clip_max=clip_max)
231        self.strong1 = strong1 or InvertibleAugmenter(*get_default_augmentations("strong", ndim=ndim), clip_max=clip_max)
232        self.strong2 = strong2 or InvertibleAugmenter(*get_default_augmentations("strong", ndim=ndim), clip_max=clip_max)
233
234    def reset_all(self):
235        self.weak.reset()
236        self.strong1.reset()
237        self.strong2.reset()

UniMatchv2Augmenters(ndim: int, weak=None, strong1=None, strong2=None, clip_max=None) View Source

222    def __init__(
223        self,
224        ndim: int,
225        weak=None,
226        strong1=None,
227        strong2=None,
228        clip_max=None,
229    ):
230        self.weak = weak or InvertibleAugmenter(*get_default_augmentations("weak", ndim=ndim), clip_max=clip_max)
231        self.strong1 = strong1 or InvertibleAugmenter(*get_default_augmentations("strong", ndim=ndim), clip_max=clip_max)
232        self.strong2 = strong2 or InvertibleAugmenter(*get_default_augmentations("strong", ndim=ndim), clip_max=clip_max)

weak

strong1

strong2

def reset_all(self): View Source

234    def reset_all(self):
235        self.weak.reset()
236        self.strong1.reset()
237        self.strong2.reset()