# -------------------------------------------------------- # Images Speak in Images: A Generalist Painter for In-Context Visual Learning (https://arxiv.org/abs/2212.02499) # Github source: https://github.com/baaivision/Painter # Copyright (c) 2022 Beijing Academy of Artificial Intelligence (BAAI) # Licensed under The MIT License [see LICENSE for details] # By Xinlong Wang, Wen Wang # Based on MAE, BEiT, detectron2, Mask2Former, bts, mmcv, mmdetetection, mmpose, MIRNet, MPRNet, and Uformer codebases # --------------------------------------------------------' import math import numbers import random import warnings from collections.abc import Sequence from typing import List, Optional, Tuple import torch from torch import Tensor import torchvision.transforms as transforms try: import accimage except ImportError: accimage = None import torchvision.transforms.functional as F from torchvision.transforms.functional import _interpolation_modes_from_int, InterpolationMode from PIL import Image, ImageFilter, ImageOps __all__ = [ "Compose", "ToTensor", "Normalize", "RandomHorizontalFlip", "RandomResizedCrop", ] class Compose(transforms.Compose): """Composes several transforms together. This transform does not support torchscript. Please, see the note below. Args: transforms (list of ``Transform`` objects): list of transforms to compose. """ def __init__(self, transforms): super().__init__(transforms) def __call__(self, img, tgt, interpolation1=None, interpolation2=None): for t in self.transforms: img, tgt = t(img, tgt, interpolation1=interpolation1, interpolation2=interpolation2) return img, tgt class ToTensor(transforms.ToTensor): """Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor. This transform does not support torchscript. Converts a PIL Image or numpy.ndarray (H x W x C) in the range [0, 255] to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0] if the PIL Image belongs to one of the modes (L, LA, P, I, F, RGB, YCbCr, RGBA, CMYK, 1) or if the numpy.ndarray has dtype = np.uint8 In the other cases, tensors are returned without scaling. .. note:: Because the input image is scaled to [0.0, 1.0], this transformation should not be used when transforming target image masks. See the `references`_ for implementing the transforms for image masks. .. _references: https://github.com/pytorch/vision/tree/main/references/segmentation """ def __init__(self) -> None: super().__init__() def __call__(self, pic1, pic2, interpolation1=None, interpolation2=None): """ Args: pic (PIL Image or numpy.ndarray): Image to be converted to tensor. Returns: Tensor: Converted image. """ return F.to_tensor(pic1), F.to_tensor(pic2) class Normalize(transforms.Normalize): """Normalize a tensor image with mean and standard deviation. This transform does not support PIL Image. Given mean: ``(mean[1],...,mean[n])`` and std: ``(std[1],..,std[n])`` for ``n`` channels, this transform will normalize each channel of the input ``torch.*Tensor`` i.e., ``output[channel] = (input[channel] - mean[channel]) / std[channel]`` .. note:: This transform acts out of place, i.e., it does not mutate the input tensor. Args: mean (sequence): Sequence of means for each channel. std (sequence): Sequence of standard deviations for each channel. inplace(bool,optional): Bool to make this operation in-place. """ def __init__(self, mean, std, inplace=False): super().__init__(mean, std, inplace) def forward(self, tensor1: Tensor, tensor2: Tensor, interpolation1=None, interpolation2=None): """ Args: tensor (Tensor): Tensor image to be normalized. Returns: Tensor: Normalized Tensor image. """ return F.normalize(tensor1, self.mean, self.std, self.inplace), F.normalize(tensor2, self.mean, self.std, self.inplace) class RandomResizedCrop(transforms.RandomResizedCrop): """Crop a random portion of image and resize it to a given size. If the image is torch Tensor, it is expected to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions A crop of the original image is made: the crop has a random area (H * W) and a random aspect ratio. This crop is finally resized to the given size. This is popularly used to train the Inception networks. Args: size (int or sequence): expected output size of the crop, for each edge. If size is an int instead of sequence like (h, w), a square output size ``(size, size)`` is made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]). .. note:: In torchscript mode size as single int is not supported, use a sequence of length 1: ``[size, ]``. scale (tuple of float): Specifies the lower and upper bounds for the random area of the crop, before resizing. The scale is defined with respect to the area of the original image. ratio (tuple of float): lower and upper bounds for the random aspect ratio of the crop, before resizing. interpolation (InterpolationMode): Desired interpolation enum defined by :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.BILINEAR``. If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` and ``InterpolationMode.BICUBIC`` are supported. For backward compatibility integer values (e.g. ``PIL.Image[.Resampling].NEAREST``) are still accepted, but deprecated since 0.13 and will be removed in 0.15. Please use InterpolationMode enum. """ def __init__( self, size, scale=(0.08, 1.0), ratio=(3.0 / 4.0, 4.0 / 3.0), interpolation=InterpolationMode.BILINEAR, ): super().__init__(size, scale=scale, ratio=ratio, interpolation=interpolation) def forward(self, img, tgt, interpolation1=None, interpolation2=None): """ Args: img (PIL Image or Tensor): Image to be cropped and resized. Returns: PIL Image or Tensor: Randomly cropped and resized image. """ i, j, h, w = self.get_params(img, self.scale, self.ratio) if interpolation1 == 'nearest': interpolation1 = InterpolationMode.NEAREST else: interpolation1 = InterpolationMode.BICUBIC if interpolation2 == 'nearest': interpolation2 = InterpolationMode.NEAREST else: interpolation2 = InterpolationMode.BICUBIC return F.resized_crop(img, i, j, h, w, self.size, interpolation1), \ F.resized_crop(tgt, i, j, h, w, self.size, interpolation2) class RandomHorizontalFlip(transforms.RandomHorizontalFlip): """Horizontally flip the given image randomly with a given probability. If the image is torch Tensor, it is expected to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions Args: p (float): probability of the image being flipped. Default value is 0.5 """ def __init__(self, p=0.5): super().__init__(p=p) def forward(self, img, tgt, interpolation1=None, interpolation2=None): """ Args: img (PIL Image or Tensor): Image to be flipped. Returns: PIL Image or Tensor: Randomly flipped image. """ if torch.rand(1) < self.p: return F.hflip(img), F.hflip(tgt) return img, tgt class RandomApply(transforms.RandomApply): """Apply randomly a list of transformations with a given probability. .. note:: In order to script the transformation, please use ``torch.nn.ModuleList`` as input instead of list/tuple of transforms as shown below: >>> transforms = transforms.RandomApply(torch.nn.ModuleList([ >>> transforms.ColorJitter(), >>> ]), p=0.3) >>> scripted_transforms = torch.jit.script(transforms) Make sure to use only scriptable transformations, i.e. that work with ``torch.Tensor``, does not require `lambda` functions or ``PIL.Image``. Args: transforms (sequence or torch.nn.Module): list of transformations p (float): probability """ def __init__(self, transforms, p=0.5): super().__init__(transforms, p=p) def forward(self, img, tgt, interpolation1=None, interpolation2=None): if self.p < torch.rand(1): return img, tgt for t in self.transforms: img, tgt = t(img, tgt) return img, tgt class ColorJitter(transforms.ColorJitter): """Randomly change the brightness, contrast, saturation and hue of an image. If the image is torch Tensor, it is expected to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions. If img is PIL Image, mode "1", "I", "F" and modes with transparency (alpha channel) are not supported. Args: brightness (float or tuple of float (min, max)): How much to jitter brightness. brightness_factor is chosen uniformly from [max(0, 1 - brightness), 1 + brightness] or the given [min, max]. Should be non negative numbers. contrast (float or tuple of float (min, max)): How much to jitter contrast. contrast_factor is chosen uniformly from [max(0, 1 - contrast), 1 + contrast] or the given [min, max]. Should be non negative numbers. saturation (float or tuple of float (min, max)): How much to jitter saturation. saturation_factor is chosen uniformly from [max(0, 1 - saturation), 1 + saturation] or the given [min, max]. Should be non negative numbers. hue (float or tuple of float (min, max)): How much to jitter hue. hue_factor is chosen uniformly from [-hue, hue] or the given [min, max]. Should have 0<= hue <= 0.5 or -0.5 <= min <= max <= 0.5. To jitter hue, the pixel values of the input image has to be non-negative for conversion to HSV space; thus it does not work if you normalize your image to an interval with negative values, or use an interpolation that generates negative values before using this function. """ def __init__(self, brightness=0, contrast=0, saturation=0, hue=0): super().__init__(brightness=brightness, contrast=contrast, saturation=saturation, hue=hue) def forward(self, img, tgt, interpolation1=None, interpolation2=None): """ Args: img (PIL Image or Tensor): Input image. Returns: PIL Image or Tensor: Color jittered image. """ fn_idx, brightness_factor, contrast_factor, saturation_factor, hue_factor = self.get_params( self.brightness, self.contrast, self.saturation, self.hue ) for fn_id in fn_idx: if fn_id == 0 and brightness_factor is not None: img = F.adjust_brightness(img, brightness_factor) elif fn_id == 1 and contrast_factor is not None: img = F.adjust_contrast(img, contrast_factor) elif fn_id == 2 and saturation_factor is not None: img = F.adjust_saturation(img, saturation_factor) elif fn_id == 3 and hue_factor is not None: img = F.adjust_hue(img, hue_factor) return img, tgt class RandomErasing(transforms.RandomErasing): """Randomly selects a rectangle region in a torch.Tensor image and erases its pixels. This transform does not support PIL Image. 'Random Erasing Data Augmentation' by Zhong et al. See https://arxiv.org/abs/1708.04896 Args: p: probability that the random erasing operation will be performed. scale: range of proportion of erased area against input image. ratio: range of aspect ratio of erased area. value: erasing value. Default is 0. If a single int, it is used to erase all pixels. If a tuple of length 3, it is used to erase R, G, B channels respectively. If a str of 'random', erasing each pixel with random values. inplace: boolean to make this transform inplace. Default set to False. Returns: Erased Image. Example: >>> transform = transforms.Compose([ >>> transforms.RandomHorizontalFlip(), >>> transforms.PILToTensor(), >>> transforms.ConvertImageDtype(torch.float), >>> transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), >>> transforms.RandomErasing(), >>> ]) """ def __init__(self, p=0.5, scale=(0.02, 0.33), ratio=(0.3, 3.3), value=0, inplace=False): super().__init__(p=p, scale=scale, ratio=ratio, value=value, inplace=inplace) def forward(self, img, tgt, interpolation1=None, interpolation2=None): """ Args: img (Tensor): Tensor image to be erased. Returns: img (Tensor): Erased Tensor image. """ if torch.rand(1) < self.p: # cast self.value to script acceptable type if isinstance(self.value, (int, float)): value = [self.value] elif isinstance(self.value, str): value = None elif isinstance(self.value, tuple): value = list(self.value) else: value = self.value if value is not None and not (len(value) in (1, img.shape[-3])): raise ValueError( "If value is a sequence, it should have either a single value or " f"{img.shape[-3]} (number of input channels)" ) x, y, h, w, v = self.get_params(img, scale=self.scale, ratio=self.ratio, value=value) return F.erase(img, x, y, h, w, v, self.inplace), tgt return img, tgt class GaussianBlur(object): """Gaussian blur augmentation from SimCLR: https://arxiv.org/abs/2002.05709""" def __init__(self, sigma=[.1, 2.]): self.sigma = sigma def __call__(self, img, tgt, interpolation1=None, interpolation2=None): sigma = random.uniform(self.sigma[0], self.sigma[1]) img = img.filter(ImageFilter.GaussianBlur(radius=sigma)) return img, tgt def __repr__(self) -> str: s = f"{self.__class__.__name__}( sigma={self.sigma})" return s