# -*- coding: utf-8 -*- # Copyright (c) Facebook, Inc. and its affiliates. import inspect import numpy as np import pprint from typing import Any, List, Optional, Tuple, Union from fvcore.transforms.transform import Transform, TransformList """ See "Data Augmentation" tutorial for an overview of the system: https://detectron2.readthedocs.io/tutorials/augmentation.html """ __all__ = [ "Augmentation", "AugmentationList", "AugInput", "TransformGen", "apply_transform_gens", "StandardAugInput", "apply_augmentations", ] def _check_img_dtype(img): assert isinstance(img, np.ndarray), "[Augmentation] Needs an numpy array, but got a {}!".format( type(img) ) assert not isinstance(img.dtype, np.integer) or ( img.dtype == np.uint8 ), "[Augmentation] Got image of type {}, use uint8 or floating points instead!".format( img.dtype ) assert img.ndim in [2, 3], img.ndim def _get_aug_input_args(aug, aug_input) -> List[Any]: """ Get the arguments to be passed to ``aug.get_transform`` from the input ``aug_input``. """ if aug.input_args is None: # Decide what attributes are needed automatically prms = list(inspect.signature(aug.get_transform).parameters.items()) # The default behavior is: if there is one parameter, then its "image" # (work automatically for majority of use cases, and also avoid BC breaking), # Otherwise, use the argument names. if len(prms) == 1: names = ("image",) else: names = [] for name, prm in prms: if prm.kind in ( inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD, ): raise TypeError( f""" \ The default implementation of `{type(aug)}.__call__` does not allow \ `{type(aug)}.get_transform` to use variable-length arguments (*args, **kwargs)! \ If arguments are unknown, reimplement `__call__` instead. \ """ ) names.append(name) aug.input_args = tuple(names) args = [] for f in aug.input_args: try: args.append(getattr(aug_input, f)) except AttributeError as e: raise AttributeError( f"{type(aug)}.get_transform needs input attribute '{f}', " f"but it is not an attribute of {type(aug_input)}!" ) from e return args class Augmentation: """ Augmentation defines (often random) policies/strategies to generate :class:`Transform` from data. It is often used for pre-processing of input data. A "policy" that generates a :class:`Transform` may, in the most general case, need arbitrary information from input data in order to determine what transforms to apply. Therefore, each :class:`Augmentation` instance defines the arguments needed by its :meth:`get_transform` method. When called with the positional arguments, the :meth:`get_transform` method executes the policy. Note that :class:`Augmentation` defines the policies to create a :class:`Transform`, but not how to execute the actual transform operations to those data. Its :meth:`__call__` method will use :meth:`AugInput.transform` to execute the transform. The returned `Transform` object is meant to describe deterministic transformation, which means it can be re-applied on associated data, e.g. the geometry of an image and its segmentation masks need to be transformed together. (If such re-application is not needed, then determinism is not a crucial requirement.) """ input_args: Optional[Tuple[str]] = None """ Stores the attribute names needed by :meth:`get_transform`, e.g. ``("image", "sem_seg")``. By default, it is just a tuple of argument names in :meth:`self.get_transform`, which often only contain "image". As long as the argument name convention is followed, there is no need for users to touch this attribute. """ def _init(self, params=None): if params: for k, v in params.items(): if k != "self" and not k.startswith("_"): setattr(self, k, v) def get_transform(self, *args) -> Transform: """ Execute the policy based on input data, and decide what transform to apply to inputs. Args: args: Any fixed-length positional arguments. By default, the name of the arguments should exist in the :class:`AugInput` to be used. Returns: Transform: Returns the deterministic transform to apply to the input. Examples: :: class MyAug: # if a policy needs to know both image and semantic segmentation def get_transform(image, sem_seg) -> T.Transform: pass tfm: Transform = MyAug().get_transform(image, sem_seg) new_image = tfm.apply_image(image) Notes: Users can freely use arbitrary new argument names in custom :meth:`get_transform` method, as long as they are available in the input data. In detectron2 we use the following convention: * image: (H,W) or (H,W,C) ndarray of type uint8 in range [0, 255], or floating point in range [0, 1] or [0, 255]. * boxes: (N,4) ndarray of float32. It represents the instance bounding boxes of N instances. Each is in XYXY format in unit of absolute coordinates. * sem_seg: (H,W) ndarray of type uint8. Each element is an integer label of pixel. We do not specify convention for other types and do not include builtin :class:`Augmentation` that uses other types in detectron2. """ raise NotImplementedError def __call__(self, aug_input) -> Transform: """ Augment the given `aug_input` **in-place**, and return the transform that's used. This method will be called to apply the augmentation. In most augmentation, it is enough to use the default implementation, which calls :meth:`get_transform` using the inputs. But a subclass can overwrite it to have more complicated logic. Args: aug_input (AugInput): an object that has attributes needed by this augmentation (defined by ``self.get_transform``). Its ``transform`` method will be called to in-place transform it. Returns: Transform: the transform that is applied on the input. """ args = _get_aug_input_args(self, aug_input) tfm = self.get_transform(*args) assert isinstance(tfm, (Transform, TransformList)), ( f"{type(self)}.get_transform must return an instance of Transform! " f"Got {type(tfm)} instead." ) aug_input.transform(tfm) return tfm def _rand_range(self, low=1.0, high=None, size=None): """ Uniform float random number between low and high. """ if high is None: low, high = 0, low if size is None: size = [] return np.random.uniform(low, high, size) def __repr__(self): """ Produce something like: "MyAugmentation(field1={self.field1}, field2={self.field2})" """ try: sig = inspect.signature(self.__init__) classname = type(self).__name__ argstr = [] for name, param in sig.parameters.items(): assert ( param.kind != param.VAR_POSITIONAL and param.kind != param.VAR_KEYWORD ), "The default __repr__ doesn't support *args or **kwargs" assert hasattr(self, name), ( "Attribute {} not found! " "Default __repr__ only works if attributes match the constructor.".format(name) ) attr = getattr(self, name) default = param.default if default is attr: continue attr_str = pprint.pformat(attr) if "\n" in attr_str: # don't show it if pformat decides to use >1 lines attr_str = "..." argstr.append("{}={}".format(name, attr_str)) return "{}({})".format(classname, ", ".join(argstr)) except AssertionError: return super().__repr__() __str__ = __repr__ class _TransformToAug(Augmentation): def __init__(self, tfm: Transform): self.tfm = tfm def get_transform(self, *args): return self.tfm def __repr__(self): return repr(self.tfm) __str__ = __repr__ def _transform_to_aug(tfm_or_aug): """ Wrap Transform into Augmentation. Private, used internally to implement augmentations. """ assert isinstance(tfm_or_aug, (Transform, Augmentation)), tfm_or_aug if isinstance(tfm_or_aug, Augmentation): return tfm_or_aug else: return _TransformToAug(tfm_or_aug) class AugmentationList(Augmentation): """ Apply a sequence of augmentations. It has ``__call__`` method to apply the augmentations. Note that :meth:`get_transform` method is impossible (will throw error if called) for :class:`AugmentationList`, because in order to apply a sequence of augmentations, the kth augmentation must be applied first, to provide inputs needed by the (k+1)th augmentation. """ def __init__(self, augs): """ Args: augs (list[Augmentation or Transform]): """ super().__init__() self.augs = [_transform_to_aug(x) for x in augs] def __call__(self, aug_input) -> TransformList: tfms = [] for x in self.augs: tfm = x(aug_input) tfms.append(tfm) return TransformList(tfms) def __repr__(self): msgs = [str(x) for x in self.augs] return "AugmentationList[{}]".format(", ".join(msgs)) __str__ = __repr__ class AugInput: """ Input that can be used with :meth:`Augmentation.__call__`. This is a standard implementation for the majority of use cases. This class provides the standard attributes **"image", "boxes", "sem_seg"** defined in :meth:`__init__` and they may be needed by different augmentations. Most augmentation policies do not need attributes beyond these three. After applying augmentations to these attributes (using :meth:`AugInput.transform`), the returned transforms can then be used to transform other data structures that users have. Examples: :: input = AugInput(image, boxes=boxes) tfms = augmentation(input) transformed_image = input.image transformed_boxes = input.boxes transformed_other_data = tfms.apply_other(other_data) An extended project that works with new data types may implement augmentation policies that need other inputs. An algorithm may need to transform inputs in a way different from the standard approach defined in this class. In those rare situations, users can implement a class similar to this class, that satify the following condition: * The input must provide access to these data in the form of attribute access (``getattr``). For example, if an :class:`Augmentation` to be applied needs "image" and "sem_seg" arguments, its input must have the attribute "image" and "sem_seg". * The input must have a ``transform(tfm: Transform) -> None`` method which in-place transforms all its attributes. """ # TODO maybe should support more builtin data types here def __init__( self, image: np.ndarray, *, boxes: Optional[np.ndarray] = None, sem_seg: Optional[np.ndarray] = None, ): """ Args: image (ndarray): (H,W) or (H,W,C) ndarray of type uint8 in range [0, 255], or floating point in range [0, 1] or [0, 255]. The meaning of C is up to users. boxes (ndarray or None): Nx4 float32 boxes in XYXY_ABS mode sem_seg (ndarray or None): HxW uint8 semantic segmentation mask. Each element is an integer label of pixel. """ _check_img_dtype(image) self.image = image self.boxes = boxes self.sem_seg = sem_seg def transform(self, tfm: Transform) -> None: """ In-place transform all attributes of this class. By "in-place", it means after calling this method, accessing an attribute such as ``self.image`` will return transformed data. """ self.image = tfm.apply_image(self.image) if self.boxes is not None: self.boxes = tfm.apply_box(self.boxes) if self.sem_seg is not None: self.sem_seg = tfm.apply_segmentation(self.sem_seg) def apply_augmentations( self, augmentations: List[Union[Augmentation, Transform]] ) -> TransformList: """ Equivalent of ``AugmentationList(augmentations)(self)`` """ return AugmentationList(augmentations)(self) def apply_augmentations(augmentations: List[Union[Transform, Augmentation]], inputs): """ Use ``T.AugmentationList(augmentations)(inputs)`` instead. """ if isinstance(inputs, np.ndarray): # handle the common case of image-only Augmentation, also for backward compatibility image_only = True inputs = AugInput(inputs) else: image_only = False tfms = inputs.apply_augmentations(augmentations) return inputs.image if image_only else inputs, tfms apply_transform_gens = apply_augmentations """ Alias for backward-compatibility. """ TransformGen = Augmentation """ Alias for Augmentation, since it is something that generates :class:`Transform`s """ StandardAugInput = AugInput """ Alias for compatibility. It's not worth the complexity to have two classes. """