from typing import Any, Dict, List, Optional, Tuple, Union, cast import torch from torch import Tensor from torch.nn.functional import pad from kornia.augmentation import random_generator as rg from kornia.augmentation._2d.geometric.base import GeometricAugmentationBase2D from kornia.augmentation.utils import _transform_input, _transform_output_shape, override_parameters from kornia.constants import Resample from kornia.geometry.transform import crop_by_indices, crop_by_transform_mat, get_perspective_transform class RandomCrop(GeometricAugmentationBase2D): r"""Crop random patches of a tensor image on a given size. .. image:: _static/img/RandomCrop.png Args: size: Desired output size (out_h, out_w) of the crop. Must be Tuple[int, int], then out_h = size[0], out_w = size[1]. padding: Optional padding on each border of the image. Default is None, i.e no padding. If a sequence of length 4 is provided, it is used to pad left, top, right, bottom borders respectively. If a sequence of length 2 is provided, it is used to pad left/right, top/bottom borders, respectively. pad_if_needed: It will pad the image if smaller than the desired size to avoid raising an exception. Since cropping is done after padding, the padding seems to be done at a random offset. fill: Pixel fill value for constant fill. Default is 0. If a tuple of length 3, it is used to fill R, G, B channels respectively. This value is only used when the padding_mode is constant. padding_mode: Type of padding. Should be: constant, reflect, replicate. resample: the interpolation mode. same_on_batch: apply the same transformation across the batch. align_corners: interpolation flag. p: probability of applying the transformation for the whole batch. keepdim: whether to keep the output shape the same as input (True) or broadcast it to the batch form (False). cropping_mode: The used algorithm to crop. ``slice`` will use advanced slicing to extract the tensor based on the sampled indices. ``resample`` will use `warp_affine` using the affine transformation to extract and resize at once. Use `slice` for efficiency, or `resample` for proper differentiability. Shape: - Input: :math:`(C, H, W)` or :math:`(B, C, H, W)`, Optional: :math:`(B, 3, 3)` - Output: :math:`(B, C, out_h, out_w)` Note: Input tensor must be float and normalized into [0, 1] for the best differentiability support. Additionally, this function accepts another transformation tensor (:math:`(B, 3, 3)`), then the applied transformation will be merged int to the input transformation tensor and returned. Examples: >>> import torch >>> _ = torch.manual_seed(0) >>> inputs = torch.arange(1*1*3*3.).view(1, 1, 3, 3) >>> aug = RandomCrop((2, 2), p=1., cropping_mode="resample") >>> out = aug(inputs) >>> out tensor([[[[3., 4.], [6., 7.]]]]) >>> aug.inverse(out, padding_mode="replicate") tensor([[[[3., 4., 4.], [3., 4., 4.], [6., 7., 7.]]]]) To apply the exact augmenation again, you may take the advantage of the previous parameter state: >>> input = torch.randn(1, 3, 32, 32) >>> aug = RandomCrop((2, 2), p=1., cropping_mode="resample") >>> (aug(input) == aug(input, params=aug._params)).all() tensor(True) """ def __init__( self, size: Tuple[int, int], padding: Optional[Union[int, Tuple[int, int], Tuple[int, int, int, int]]] = None, pad_if_needed: Optional[bool] = False, fill: int = 0, padding_mode: str = "constant", resample: Union[str, int, Resample] = Resample.BILINEAR.name, same_on_batch: bool = False, align_corners: bool = True, p: float = 1.0, keepdim: bool = False, cropping_mode: str = "slice", return_transform: Optional[bool] = None, ) -> None: # Since PyTorch does not support ragged tensor. So cropping function happens batch-wisely. super().__init__( p=1.0, return_transform=return_transform, same_on_batch=same_on_batch, p_batch=p, keepdim=keepdim ) self._param_generator = cast(rg.CropGenerator, rg.CropGenerator(size)) self.flags = dict( size=size, padding=padding, pad_if_needed=pad_if_needed, fill=fill, padding_mode=padding_mode, resample=Resample.get(resample), align_corners=align_corners, cropping_mode=cropping_mode, ) def compute_padding(self, shape: torch.Size, flags: Optional[Dict[str, Any]] = None) -> List[int]: flags = self.flags if flags is None else flags if len(shape) != 4: raise AssertionError(f"Expected BCHW. Got {shape}.") padding = [0, 0, 0, 0] if flags["padding"] is not None: if isinstance(flags["padding"], int): padding = [flags["padding"]] * 4 elif isinstance(flags["padding"], tuple) and len(flags["padding"]) == 2: padding = [flags["padding"][1], flags["padding"][1], flags["padding"][0], flags["padding"][0]] elif isinstance(flags["padding"], tuple) and len(flags["padding"]) == 4: padding = [flags["padding"][3], flags["padding"][2], flags["padding"][1], flags["padding"][0]] else: raise RuntimeError(f"Expect `padding` to be a scalar, or length 2/4 list. Got {flags['padding']}.") if flags["pad_if_needed"] and shape[-2] < flags["size"][0]: padding = [0, 0, (flags["size"][0] - shape[-2]), flags["size"][0] - shape[-2]] if flags["pad_if_needed"] and shape[-1] < flags["size"][1]: padding = [flags["size"][1] - shape[-1], flags["size"][1] - shape[-1], 0, 0] return padding def precrop_padding( self, input: Tensor, padding: List[int] = None, flags: Optional[Dict[str, Any]] = None ) -> Tensor: flags = self.flags if flags is None else flags if padding is None: padding = self.compute_padding(input.shape) input = pad(input, padding, value=flags["fill"], mode=flags["padding_mode"]) return input def compute_transformation(self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any]) -> Tensor: if flags["cropping_mode"] in ("resample", "slice"): transform: Tensor = get_perspective_transform(params["src"].to(input), params["dst"].to(input)) return transform raise NotImplementedError(f"Not supported type: {flags['cropping_mode']}.") def apply_transform( self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None ) -> Tensor: flags = self.flags if flags is None else flags if flags["cropping_mode"] == "resample": # uses bilinear interpolation to crop transform = cast(Tensor, transform) # Fit the arg to F.pad if flags['padding_mode'] == "constant": padding_mode = "zeros" elif flags['padding_mode'] == "replicate": padding_mode = "border" elif flags['padding_mode'] == "reflect": padding_mode = "reflection" else: padding_mode = flags['padding_mode'] return crop_by_transform_mat( input, transform, flags["size"], mode=flags["resample"].name.lower(), padding_mode=padding_mode, align_corners=flags["align_corners"], ) if flags["cropping_mode"] == "slice": # uses advanced slicing to crop return crop_by_indices(input, params["src"], flags["size"]) raise NotImplementedError(f"Not supported type: {flags['cropping_mode']}.") def inverse_transform( self, input: Tensor, flags: Dict[str, Any], transform: Optional[Tensor] = None, size: Optional[Tuple[int, int]] = None, ) -> Tensor: if flags["cropping_mode"] != "resample": raise NotImplementedError( f"`inverse` is only applicable for resample cropping mode. Got {flags['cropping_mode']}." ) size = cast(Tuple[int, int], size) transform = cast(Tensor, transform) # Fit the arg to F.pad if flags['padding_mode'] == "constant": padding_mode = "zeros" elif flags['padding_mode'] == "replicate": padding_mode = "border" elif flags['padding_mode'] == "reflect": padding_mode = "reflection" else: padding_mode = flags['padding_mode'] return crop_by_transform_mat( input, transform[:, :2, :], size, flags["resample"].name.lower(), padding_mode=padding_mode, align_corners=flags["align_corners"], ) def inverse( self, input: Tensor, params: Optional[Dict[str, Tensor]] = None, size: Optional[Tuple[int, int]] = None, **kwargs, ) -> Tensor: out = super().inverse(input, params, size, **kwargs) if params is None: params = self._params if "padding_size" in params: padding_size = params["padding_size"].unique(dim=0).cpu().squeeze().numpy().tolist() padding_size = [-padding_size[0], -padding_size[1], -padding_size[2], -padding_size[3]] else: padding_size = [0, 0, 0, 0] return self.precrop_padding(out, padding_size) def forward_parameters_precrop(self, batch_shape) -> Dict[str, Tensor]: input_pad = self.compute_padding(batch_shape) batch_shape_new = ( *batch_shape[:2], batch_shape[2] + input_pad[2] + input_pad[3], # original height + top + bottom padding batch_shape[3] + input_pad[0] + input_pad[1], # original width + left + right padding ) padding_size = torch.tensor(tuple(input_pad), dtype=torch.long).expand(batch_shape[0], -1) _params = super().forward_parameters(batch_shape_new) _params.update({"padding_size": padding_size}) return _params def forward(self, input: Tensor, params: Optional[Dict[str, Tensor]] = None, **kwargs) -> Tensor: padding_size = params.get("padding_size") if params else None if padding_size is not None: input_pad = padding_size.unique(dim=0).cpu().squeeze().numpy().tolist() else: input_pad = None flags = override_parameters(self.flags, kwargs, in_place=False) if isinstance(input, (tuple, list)): ori_shape = input[0].shape input_temp = _transform_input(input[0]) input_pad = self.compute_padding(input[0].shape, flags) if input_pad is None else input_pad _input = (self.precrop_padding(input_temp, input_pad, flags), input[1]) _input = _transform_output_shape(_input, ori_shape) if self.keepdim else _input # type:ignore else: input = cast(Tensor, input) ori_shape = input.shape input_temp = _transform_input(input) input_pad = self.compute_padding(input_temp.shape, flags) if input_pad is None else input_pad _input = self.precrop_padding(input_temp, input_pad, flags) # type: ignore _input = _transform_output_shape(_input, ori_shape) if self.keepdim else _input # type:ignore if params is not None: params, flags = self._process_kwargs_to_params_and_flags(params, self.flags, **kwargs) out = super().forward(_input, params, **kwargs) # type:ignore # Update the actual input size for inverse if "padding_size" not in self._params: _padding_size = torch.tensor(tuple(input_pad), device=input_temp.device, dtype=torch.long).expand( input_temp.size(0), -1 ) self._params.update({"padding_size": _padding_size}) if not self._params["batch_prob"].all(): # undo the pre-crop if nothing happened. if isinstance(out, tuple) and isinstance(input, tuple): return input[0], out[1] if isinstance(out, tuple) and not isinstance(input, tuple): return input, out[1] return input return out