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 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_transform(
self, input: Tensor, params: Dict[str, Tensor], transform: Optional[Tensor] = None
) -> Tensor:
# sample a kernel size
kernel_size_list: List[int] = params["ksize_factor"].tolist()
idx: int = cast(int, torch.randint(len(kernel_size_list), (1,)).item())
return motion_blur(
input,
kernel_size=kernel_size_list[idx],
angle=params["angle_factor"],
direction=params["direction_factor"],
border_type=self.flags["border_type"].name.lower(),
mode=self.flags["resample"].name.lower(),
)
def test_against_functional(self, input_shape):
input = torch.randn(*input_shape)
f = RandomMotionBlur(kernel_size=(3, 5), angle=(10, 30), direction=0.5, p=1.0)
output = f(input)
expected = motion_blur(
input,
f._params['ksize_factor'].unique().item(),
f._params['angle_factor'],
f._params['direction_factor'],
f.border_type.name.lower(),
)
assert_allclose(output, expected, rtol=1e-4, atol=1e-4)
