jawahar-konathala/Tryon2

1# -*- coding: utf-8 -*-2# Copyright (c) Facebook, Inc. and its affiliates.3 4import inspect5import numpy as np6import pprint7from typing import Any, List, Optional, Tuple, Union8from fvcore.transforms.transform import Transform, TransformList9 10"""11See "Data Augmentation" tutorial for an overview of the system:12https://detectron2.readthedocs.io/tutorials/augmentation.html13"""14 15 16__all__ = [17    "Augmentation",18    "AugmentationList",19    "AugInput",20    "TransformGen",21    "apply_transform_gens",22    "StandardAugInput",23    "apply_augmentations",24]25 26 27def _check_img_dtype(img):28    assert isinstance(img, np.ndarray), "[Augmentation] Needs an numpy array, but got a {}!".format(29        type(img)30    )31    assert not isinstance(img.dtype, np.integer) or (32        img.dtype == np.uint833    ), "[Augmentation] Got image of type {}, use uint8 or floating points instead!".format(34        img.dtype35    )36    assert img.ndim in [2, 3], img.ndim37 38 39def _get_aug_input_args(aug, aug_input) -> List[Any]:40    """41    Get the arguments to be passed to ``aug.get_transform`` from the input ``aug_input``.42    """43    if aug.input_args is None:44        # Decide what attributes are needed automatically45        prms = list(inspect.signature(aug.get_transform).parameters.items())46        # The default behavior is: if there is one parameter, then its "image"47        # (work automatically for majority of use cases, and also avoid BC breaking),48        # Otherwise, use the argument names.49        if len(prms) == 1:50            names = ("image",)51        else:52            names = []53            for name, prm in prms:54                if prm.kind in (55                    inspect.Parameter.VAR_POSITIONAL,56                    inspect.Parameter.VAR_KEYWORD,57                ):58                    raise TypeError(59                        f""" \60The default implementation of `{type(aug)}.__call__` does not allow \61`{type(aug)}.get_transform` to use variable-length arguments (*args, **kwargs)! \62If arguments are unknown, reimplement `__call__` instead. \63"""64                    )65                names.append(name)66        aug.input_args = tuple(names)67 68    args = []69    for f in aug.input_args:70        try:71            args.append(getattr(aug_input, f))72        except AttributeError as e:73            raise AttributeError(74                f"{type(aug)}.get_transform needs input attribute '{f}', "75                f"but it is not an attribute of {type(aug_input)}!"76            ) from e77    return args78 79 80class Augmentation:81    """82    Augmentation defines (often random) policies/strategies to generate :class:`Transform`83    from data. It is often used for pre-processing of input data.84 85    A "policy" that generates a :class:`Transform` may, in the most general case,86    need arbitrary information from input data in order to determine what transforms87    to apply. Therefore, each :class:`Augmentation` instance defines the arguments88    needed by its :meth:`get_transform` method. When called with the positional arguments,89    the :meth:`get_transform` method executes the policy.90 91    Note that :class:`Augmentation` defines the policies to create a :class:`Transform`,92    but not how to execute the actual transform operations to those data.93    Its :meth:`__call__` method will use :meth:`AugInput.transform` to execute the transform.94 95    The returned `Transform` object is meant to describe deterministic transformation, which means96    it can be re-applied on associated data, e.g. the geometry of an image and its segmentation97    masks need to be transformed together.98    (If such re-application is not needed, then determinism is not a crucial requirement.)99    """100 101    input_args: Optional[Tuple[str]] = None102    """103    Stores the attribute names needed by :meth:`get_transform`, e.g.  ``("image", "sem_seg")``.104    By default, it is just a tuple of argument names in :meth:`self.get_transform`, which often only105    contain "image". As long as the argument name convention is followed, there is no need for106    users to touch this attribute.107    """108 109    def _init(self, params=None):110        if params:111            for k, v in params.items():112                if k != "self" and not k.startswith("_"):113                    setattr(self, k, v)114 115    def get_transform(self, *args) -> Transform:116        """117        Execute the policy based on input data, and decide what transform to apply to inputs.118 119        Args:120            args: Any fixed-length positional arguments. By default, the name of the arguments121                should exist in the :class:`AugInput` to be used.122 123        Returns:124            Transform: Returns the deterministic transform to apply to the input.125 126        Examples:127        ::128            class MyAug:129                # if a policy needs to know both image and semantic segmentation130                def get_transform(image, sem_seg) -> T.Transform:131                    pass132            tfm: Transform = MyAug().get_transform(image, sem_seg)133            new_image = tfm.apply_image(image)134 135        Notes:136            Users can freely use arbitrary new argument names in custom137            :meth:`get_transform` method, as long as they are available in the138            input data. In detectron2 we use the following convention:139 140            * image: (H,W) or (H,W,C) ndarray of type uint8 in range [0, 255], or141              floating point in range [0, 1] or [0, 255].142            * boxes: (N,4) ndarray of float32. It represents the instance bounding boxes143              of N instances. Each is in XYXY format in unit of absolute coordinates.144            * sem_seg: (H,W) ndarray of type uint8. Each element is an integer label of pixel.145 146            We do not specify convention for other types and do not include builtin147            :class:`Augmentation` that uses other types in detectron2.148        """149        raise NotImplementedError150 151    def __call__(self, aug_input) -> Transform:152        """153        Augment the given `aug_input` **in-place**, and return the transform that's used.154 155        This method will be called to apply the augmentation. In most augmentation, it156        is enough to use the default implementation, which calls :meth:`get_transform`157        using the inputs. But a subclass can overwrite it to have more complicated logic.158 159        Args:160            aug_input (AugInput): an object that has attributes needed by this augmentation161                (defined by ``self.get_transform``). Its ``transform`` method will be called162                to in-place transform it.163 164        Returns:165            Transform: the transform that is applied on the input.166        """167        args = _get_aug_input_args(self, aug_input)168        tfm = self.get_transform(*args)169        assert isinstance(tfm, (Transform, TransformList)), (170            f"{type(self)}.get_transform must return an instance of Transform! "171            f"Got {type(tfm)} instead."172        )173        aug_input.transform(tfm)174        return tfm175 176    def _rand_range(self, low=1.0, high=None, size=None):177        """178        Uniform float random number between low and high.179        """180        if high is None:181            low, high = 0, low182        if size is None:183            size = []184        return np.random.uniform(low, high, size)185 186    def __repr__(self):187        """188        Produce something like:189        "MyAugmentation(field1={self.field1}, field2={self.field2})"190        """191        try:192            sig = inspect.signature(self.__init__)193            classname = type(self).__name__194            argstr = []195            for name, param in sig.parameters.items():196                assert (197                    param.kind != param.VAR_POSITIONAL and param.kind != param.VAR_KEYWORD198                ), "The default __repr__ doesn't support *args or **kwargs"199                assert hasattr(self, name), (200                    "Attribute {} not found! "201                    "Default __repr__ only works if attributes match the constructor.".format(name)202                )203                attr = getattr(self, name)204                default = param.default205                if default is attr:206                    continue207                attr_str = pprint.pformat(attr)208                if "\n" in attr_str:209                    # don't show it if pformat decides to use >1 lines210                    attr_str = "..."211                argstr.append("{}={}".format(name, attr_str))212            return "{}({})".format(classname, ", ".join(argstr))213        except AssertionError:214            return super().__repr__()215 216    __str__ = __repr__217 218 219class _TransformToAug(Augmentation):220    def __init__(self, tfm: Transform):221        self.tfm = tfm222 223    def get_transform(self, *args):224        return self.tfm225 226    def __repr__(self):227        return repr(self.tfm)228 229    __str__ = __repr__230 231 232def _transform_to_aug(tfm_or_aug):233    """234    Wrap Transform into Augmentation.235    Private, used internally to implement augmentations.236    """237    assert isinstance(tfm_or_aug, (Transform, Augmentation)), tfm_or_aug238    if isinstance(tfm_or_aug, Augmentation):239        return tfm_or_aug240    else:241        return _TransformToAug(tfm_or_aug)242 243 244class AugmentationList(Augmentation):245    """246    Apply a sequence of augmentations.247 248    It has ``__call__`` method to apply the augmentations.249 250    Note that :meth:`get_transform` method is impossible (will throw error if called)251    for :class:`AugmentationList`, because in order to apply a sequence of augmentations,252    the kth augmentation must be applied first, to provide inputs needed by the (k+1)th253    augmentation.254    """255 256    def __init__(self, augs):257        """258        Args:259            augs (list[Augmentation or Transform]):260        """261        super().__init__()262        self.augs = [_transform_to_aug(x) for x in augs]263 264    def __call__(self, aug_input) -> TransformList:265        tfms = []266        for x in self.augs:267            tfm = x(aug_input)268            tfms.append(tfm)269        return TransformList(tfms)270 271    def __repr__(self):272        msgs = [str(x) for x in self.augs]273        return "AugmentationList[{}]".format(", ".join(msgs))274 275    __str__ = __repr__276 277 278class AugInput:279    """280    Input that can be used with :meth:`Augmentation.__call__`.281    This is a standard implementation for the majority of use cases.282    This class provides the standard attributes **"image", "boxes", "sem_seg"**283    defined in :meth:`__init__` and they may be needed by different augmentations.284    Most augmentation policies do not need attributes beyond these three.285 286    After applying augmentations to these attributes (using :meth:`AugInput.transform`),287    the returned transforms can then be used to transform other data structures that users have.288 289    Examples:290    ::291        input = AugInput(image, boxes=boxes)292        tfms = augmentation(input)293        transformed_image = input.image294        transformed_boxes = input.boxes295        transformed_other_data = tfms.apply_other(other_data)296 297    An extended project that works with new data types may implement augmentation policies298    that need other inputs. An algorithm may need to transform inputs in a way different299    from the standard approach defined in this class. In those rare situations, users can300    implement a class similar to this class, that satify the following condition:301 302    * The input must provide access to these data in the form of attribute access303      (``getattr``).  For example, if an :class:`Augmentation` to be applied needs "image"304      and "sem_seg" arguments, its input must have the attribute "image" and "sem_seg".305    * The input must have a ``transform(tfm: Transform) -> None`` method which306      in-place transforms all its attributes.307    """308 309    # TODO maybe should support more builtin data types here310    def __init__(311        self,312        image: np.ndarray,313        *,314        boxes: Optional[np.ndarray] = None,315        sem_seg: Optional[np.ndarray] = None,316    ):317        """318        Args:319            image (ndarray): (H,W) or (H,W,C) ndarray of type uint8 in range [0, 255], or320                floating point in range [0, 1] or [0, 255]. The meaning of C is up321                to users.322            boxes (ndarray or None): Nx4 float32 boxes in XYXY_ABS mode323            sem_seg (ndarray or None): HxW uint8 semantic segmentation mask. Each element324                is an integer label of pixel.325        """326        _check_img_dtype(image)327        self.image = image328        self.boxes = boxes329        self.sem_seg = sem_seg330 331    def transform(self, tfm: Transform) -> None:332        """333        In-place transform all attributes of this class.334 335        By "in-place", it means after calling this method, accessing an attribute such336        as ``self.image`` will return transformed data.337        """338        self.image = tfm.apply_image(self.image)339        if self.boxes is not None:340            self.boxes = tfm.apply_box(self.boxes)341        if self.sem_seg is not None:342            self.sem_seg = tfm.apply_segmentation(self.sem_seg)343 344    def apply_augmentations(345        self, augmentations: List[Union[Augmentation, Transform]]346    ) -> TransformList:347        """348        Equivalent of ``AugmentationList(augmentations)(self)``349        """350        return AugmentationList(augmentations)(self)351 352 353def apply_augmentations(augmentations: List[Union[Transform, Augmentation]], inputs):354    """355    Use ``T.AugmentationList(augmentations)(inputs)`` instead.356    """357    if isinstance(inputs, np.ndarray):358        # handle the common case of image-only Augmentation, also for backward compatibility359        image_only = True360        inputs = AugInput(inputs)361    else:362        image_only = False363    tfms = inputs.apply_augmentations(augmentations)364    return inputs.image if image_only else inputs, tfms365 366 367apply_transform_gens = apply_augmentations368"""369Alias for backward-compatibility.370"""371 372TransformGen = Augmentation373"""374Alias for Augmentation, since it is something that generates :class:`Transform`s375"""376 377StandardAugInput = AugInput378"""379Alias for compatibility. It's not worth the complexity to have two classes.380"""381