def apply_motion_blur(input: torch.Tensor, params: Dict[str, torch.Tensor],
flags: Dict[str, torch.Tensor]) -> torch.Tensor:
r"""Perform motion blur on an image.
The input image is expected to be in the range of [0, 1].
Args:
input (torch.Tensor): Tensor to be transformed with shape (H, W), (C, H, W), (B, C, H, W).
params (Dict[str, torch.Tensor]):
- params['ksize_factor']: motion kernel width and height (odd and positive).
- params['angle_factor']: angle of the motion blur in degrees (anti-clockwise rotation).
- params['direction_factor']: forward/backward direction of the motion blur.
Lower values towards -1.0 will point the motion blur towards the back (with
angle provided via angle), while higher values towards 1.0 will point the motion
blur forward. A value of 0.0 leads to a uniformly (but still angled) motion blur.
flags (Dict[str, torch.Tensor]):
- flags['border_type']: the padding mode to be applied before convolving.
CONSTANT = 0, REFLECT = 1, REPLICATE = 2, CIRCULAR = 3. Default: BorderType.CONSTANT.
Returns:
torch.Tensor: Adjusted image with the shape as the input (\*, C, H, W).
"""
kernel_size: int = cast(int, params['ksize_factor'].unique().item())
angle = params['angle_factor']
direction = params['direction_factor']
border_type: str = cast(
str,
BorderType(flags['border_type'].item()).name.lower())
mode: str = cast(str, Resample(flags['interpolation'].item()).name.lower())
return motion_blur(input, kernel_size, angle, direction, border_type, mode)
def apply_motion_blur(input: torch.Tensor,
params: Dict[str, torch.Tensor]) -> torch.Tensor:
r"""Perform motion blur on an image
The input image is expected to be in the range of [0, 1].
Args:
input (torch.Tensor): Tensor to be transformed with shape (H, W), (C, H, W), (B, C, H, W).
params (Dict[str, torch.Tensor]):
- params['ksize_factor']: motion kernel width and height (odd and positive).
- params['angle_factor']: angle of the motion blur in degrees (anti-clockwise rotation).
- params['direction_factor']: forward/backward direction of the motion blur.
Lower values towards -1.0 will point the motion blur towards the back (with
angle provided via angle), while higher values towards 1.0 will point the motion
blur forward. A value of 0.0 leads to a uniformly (but still angled) motion blur.
- params['border_type']: the padding mode to be applied before convolving.
CONSTANT = 0, REFLECT = 1, REPLICATE = 2, CIRCULAR = 3. Default: BorderType.CONSTANT.
Returns:
torch.Tensor: Adjusted image with the shape as the inpute (\*, C, H, W).
"""
input = _transform_input(input)
_validate_input_dtype(
input, accepted_dtypes=[torch.float16, torch.float32, torch.float64])
kernel_size: int = cast(int, params['ksize_factor'].item())
# TODO: this params should be at some point, learnable tensors
angle: float = cast(float, params['angle_factor'].item())
direction: float = cast(float, params['direction_factor'].item())
border_type: str = cast(
str,
BorderType(params['border_type'].item()).name.lower())
return motion_blur(input, kernel_size, angle, direction, border_type)
def __init__(
self,
kernel_size: Union[int, Tuple[int, int]],
angle: Union[
torch.Tensor,
float,
Tuple[float, float, float],
Tuple[Tuple[float, float], Tuple[float, float], Tuple[float, float]],
],
direction: Union[torch.Tensor, float, Tuple[float, float]],
border_type: Union[int, str, BorderType] = BorderType.CONSTANT.name,
resample: Union[str, int, Resample] = Resample.NEAREST.name,
return_transform: bool = False,
same_on_batch: bool = False,
p: float = 0.5,
keepdim: bool = False,
) -> None:
super().__init__(
p=p, return_transform=return_transform, same_on_batch=same_on_batch, p_batch=1.0, keepdim=keepdim
)
self.flags = dict(
border_type=BorderType.get(border_type),
resample=Resample.get(resample),
)
self._param_generator = cast(rg.MotionBlurGenerator3D, rg.MotionBlurGenerator3D(kernel_size, angle, direction))
def motion_blur_param_generator(
batch_size: int,
kernel_size: Union[int, Tuple[int, int]],
angle: UnionFloat,
direction: UnionFloat,
border_type: Union[int, str, BorderType] = BorderType.CONSTANT.name,
same_on_batch: bool = True) -> Dict[str, torch.Tensor]:
angle_bound: torch.Tensor = _check_and_bound(angle, 'angle', center=0.)
direction_bound: torch.Tensor = _check_and_bound(direction,
'direction',
center=0.,
bounds=(-1, 1))
if isinstance(kernel_size, int):
ksize_factor = torch.tensor([kernel_size] * batch_size)
elif isinstance(kernel_size, tuple):
ksize_x, ksize_y = kernel_size
ksize_factor = _adapted_uniform(
(batch_size, ), ksize_x // 2, ksize_y // 2,
same_on_batch).int() * 2 + 1
else:
raise TypeError(f"Unsupported type: {type(kernel_size)}")
angle_factor = _adapted_uniform((batch_size, ), angle_bound[0],
angle_bound[1], same_on_batch)
direction_factor = _adapted_uniform((batch_size, ), direction_bound[0],
direction_bound[1], same_on_batch)
return dict(ksize_factor=ksize_factor,
angle_factor=angle_factor,
direction_factor=direction_factor,
border_type=torch.tensor(BorderType.get(border_type).value))
def random_motion_blur_generator(
batch_size: int,
kernel_size: Union[int, Tuple[int, int]],
angle: torch.Tensor,
direction: torch.Tensor,
border_type: Union[int, str, BorderType] = BorderType.CONSTANT.name,
same_on_batch: bool = True) -> Dict[str, torch.Tensor]:
_joint_range_check(angle, 'angle')
_joint_range_check(direction, 'direction')
if isinstance(kernel_size, int):
ksize_factor = torch.tensor([kernel_size] * batch_size)
elif isinstance(kernel_size, tuple):
ksize_x, ksize_y = kernel_size
ksize_factor = _adapted_uniform(
(batch_size, ), ksize_x // 2, ksize_y // 2,
same_on_batch).int() * 2 + 1
else:
raise TypeError(f"Unsupported type: {type(kernel_size)}")
angle_factor = _adapted_uniform((batch_size, ), angle[0], angle[1],
same_on_batch)
direction_factor = _adapted_uniform((batch_size, ), direction[0],
direction[1], same_on_batch)
return dict(ksize_factor=ksize_factor,
angle_factor=angle_factor,
direction_factor=direction_factor,
border_type=torch.tensor(BorderType.get(border_type).value))
def __init__(self,
kernel_size: Union[int, Tuple[int, int]],
angle: Union[torch.Tensor, float, Tuple[float, float, float],
Tuple[Tuple[float, float], Tuple[float, float],
Tuple[float, float]]],
direction: Union[torch.Tensor, float, Tuple[float, float]],
border_type: Union[int, str,
BorderType] = BorderType.CONSTANT.name,
return_transform: bool = False,
same_on_batch: bool = False,
p: float = 0.5) -> None:
super(RandomMotionBlur3D,
self).__init__(p=p,
return_transform=return_transform,
same_on_batch=same_on_batch,
p_batch=1.)
self.kernel_size: Union[int, Tuple[int, int]] = kernel_size
self.angle: torch.Tensor = _tuple_range_reader(angle, 3)
direction = \
cast(torch.Tensor, direction) if isinstance(direction, torch.Tensor) else torch.tensor(direction)
self.direction = _range_bound(direction,
'direction',
center=0.,
bounds=(-1, 1))
self.border_type = BorderType.get(border_type)
self.flags: Dict[str, torch.Tensor] = {
"border_type": torch.tensor(self.border_type.value)
}
def __init__(
self,
kernel_size: Union[int, Tuple[int, int]],
angle: Union[
torch.Tensor,
float,
Tuple[float, float, float],
Tuple[Tuple[float, float], Tuple[float, float], Tuple[float, float]],
],
direction: Union[torch.Tensor, float, Tuple[float, float]],
border_type: Union[int, str, BorderType] = BorderType.CONSTANT.name,
resample: Union[str, int, Resample] = Resample.NEAREST.name,
return_transform: bool = False,
same_on_batch: bool = False,
p: float = 0.5,
keepdim: bool = False,
) -> None:
super(RandomMotionBlur3D, self).__init__(
p=p, return_transform=return_transform, same_on_batch=same_on_batch, p_batch=1.0, keepdim=keepdim
)
self._device, self._dtype = _extract_device_dtype([angle, direction])
self.kernel_size: Union[int, Tuple[int, int]] = kernel_size
self.angle = angle
self.direction = direction
self.resample = Resample.get(resample)
self.border_type = BorderType.get(border_type)
self.flags: Dict[str, torch.Tensor] = {
"border_type": torch.tensor(self.border_type.value),
"interpolation": torch.tensor(self.resample.value),
}
def __init__(
self, kernel_size: Union[int, Tuple[int, int]],
angle: Union[float, Tuple[float, float]],
direction: Union[float, Tuple[float, float]],
border_type: Union[int, str, BorderType] = BorderType.CONSTANT.name,
return_transform: bool = False
) -> None:
super(RandomMotionBlur, self).__init__(return_transform)
self.kernel_size: Union[int, Tuple[int, int]] = kernel_size
self.angle: Union[float, Tuple[float, float]] = angle
self.direction: Union[float, Tuple[float, float]] = direction
self.border_type: BorderType = BorderType.get(border_type)
def __init__(
self,
kernel_size: Tuple[int, int],
sigma: Tuple[float, float],
border_type: str = "reflect",
same_on_batch: bool = False,
p: float = 0.5,
keepdim: bool = False,
return_transform: Optional[bool] = None,
) -> None:
super().__init__(p=p,
return_transform=return_transform,
same_on_batch=same_on_batch,
p_batch=1.0,
keepdim=keepdim)
self.flags = dict(kernel_size=kernel_size,
sigma=sigma,
border_type=BorderType.get(border_type))
