Augmentation pipeline that supports deterministic replay. Parameters ---------- All parameters are identical to the original NVIDIA implementation. Two **extra** keyword arguments are recognised by ``__call__``: ``params`` and ``return_params``. See the docstring of ``__call__
| 141 | # Augmentation pipeline with deterministic replay capability. |
| 142 | |
| 143 | class AugmentPipe: |
| 144 | """Augmentation pipeline that supports deterministic replay. |
| 145 | |
| 146 | Parameters |
| 147 | ---------- |
| 148 | All parameters are identical to the original NVIDIA implementation. Two |
| 149 | **extra** keyword arguments are recognised by ``__call__``: |
| 150 | ``params`` and ``return_params``. See the docstring of ``__call__`` for |
| 151 | details. |
| 152 | """ |
| 153 | |
| 154 | # Original constructor signature is kept intact. |
| 155 | def __init__(self, p=1, |
| 156 | xflip=0, yflip=0, rotate_int=0, translate_int=0, translate_int_max=0.125, |
| 157 | scale=0, rotate_frac=0, aniso=0, translate_frac=0, scale_std=0.2, rotate_frac_max=1, aniso_std=0.2, aniso_rotate_prob=0.5, translate_frac_std=0.125, |
| 158 | brightness=0, contrast=0, lumaflip=0, hue=0, saturation=0, brightness_std=0.2, contrast_std=0.5, hue_max=1, saturation_std=1): |
| 159 | |
| 160 | super().__init__() |
| 161 | self.p = float(p) # Overall multiplier. |
| 162 | |
| 163 | # Pixel blitting. |
| 164 | self.xflip = float(xflip) |
| 165 | self.yflip = float(yflip) |
| 166 | self.rotate_int = float(rotate_int) |
| 167 | self.translate_int = float(translate_int) |
| 168 | self.translate_int_max = float(translate_int_max) |
| 169 | |
| 170 | # Geometric transformations. |
| 171 | self.scale = float(scale) |
| 172 | self.rotate_frac = float(rotate_frac) |
| 173 | self.aniso = float(aniso) |
| 174 | self.translate_frac = float(translate_frac) |
| 175 | self.scale_std = float(scale_std) |
| 176 | self.rotate_frac_max = float(rotate_frac_max) |
| 177 | self.aniso_std = float(aniso_std) |
| 178 | self.aniso_rotate_prob = float(aniso_rotate_prob) |
| 179 | self.translate_frac_std = float(translate_frac_std) |
| 180 | |
| 181 | # Color transformations. |
| 182 | self.brightness = float(brightness) |
| 183 | self.contrast = float(contrast) |
| 184 | self.lumaflip = float(lumaflip) |
| 185 | self.hue = float(hue) |
| 186 | self.saturation = float(saturation) |
| 187 | self.brightness_std = float(brightness_std) |
| 188 | self.contrast_std = float(contrast_std) |
| 189 | self.hue_max = float(hue_max) |
| 190 | self.saturation_std = float(saturation_std) |
| 191 | |
| 192 | # --------------------------------------------------------------------- |
| 193 | # Forward pass with optional deterministic parameters. |
| 194 | # --------------------------------------------------------------------- |
| 195 | def __call__(self, images: torch.Tensor, *, params: Optional[Dict[str, torch.Tensor]] = None, |
| 196 | return_params: bool = False) -> Tuple[torch.Tensor, torch.Tensor, Optional[Dict[str, torch.Tensor]]]: |
| 197 | """Apply augmentations. |
| 198 | |
| 199 | Parameters |
| 200 | ---------- |