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))