def _add_noise(self, img: Union[torch.Tensor, np.ndarray], mean: float, std: float):
im_shape = img.shape
_std = self.R.uniform(0, std) if self.sample_std else std
self._noise1 = self.R.normal(mean, _std, size=im_shape)
self._noise2 = self.R.normal(mean, _std, size=im_shape)
if self._noise1 is None or self._noise2 is None:
raise AssertionError
dtype = dtype_torch_to_numpy(img.dtype) if isinstance(img, torch.Tensor) else img.dtype
return np.sqrt((img + self._noise1.astype(dtype)) ** 2 + self._noise2.astype(dtype) ** 2)
def __call__(self, img: Union[torch.Tensor, np.ndarray]) -> Union[torch.Tensor, np.ndarray]:
"""
Apply the transform to `img`.
"""
self.randomize(img.shape)
assert self._noise is not None
if not self._do_transform:
return img
dtype = dtype_torch_to_numpy(img.dtype) if isinstance(img, torch.Tensor) else img.dtype
return img + self._noise.astype(dtype)
def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]:
d = dict(data)
image_shape = d[self.keys[0]].shape # image shape from the first data key
self.randomize(image_shape)
assert self._noise is not None
if not self._do_transform:
return d
for key in self.keys:
dtype = dtype_torch_to_numpy(d[key].dtype) if isinstance(d[key], torch.Tensor) else d[key].dtype
d[key] = d[key] + self._noise.astype(dtype)
return d
def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]:
d = dict(data)
image_shape = d[self.keys[0]].shape # image shape from the first data key
self.randomize(image_shape)
if len(self._noise) != len(self.keys):
raise AssertionError
if not self._do_transform:
return d
for noise, key in zip(self._noise, self.keys):
dtype = dtype_torch_to_numpy(d[key].dtype) if isinstance(d[key], torch.Tensor) else d[key].dtype
d[key] = d[key] + noise.astype(dtype)
return d
