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)