Note
You are reading the documentation for MMSelfSup 0.x, which will soon be deprecated by the end of 2022. We recommend you upgrade to MMSelfSup 1.0.0rc versions to enjoy fruitful new features and better performance brought by OpenMMLab 2.0. Check out the changelog, code and documentation of MMSelfSup 1.0.0rc for more details.
Source code for mmselfsup.datasets.pipelines.transforms
# Copyright (c) OpenMMLab. All rights reserved.
import inspect
import math
import random
import warnings
from typing import Optional, Sequence, Tuple, Union
import numpy as np
import torch
import torchvision.transforms.functional as F
from mmcv.utils import build_from_cfg
from PIL import Image, ImageFilter
from timm.data import create_transform
from torchvision import transforms as _transforms
from ..builder import PIPELINES
# register all existing transforms in torchvision
_EXCLUDED_TRANSFORMS = ['GaussianBlur']
for m in inspect.getmembers(_transforms, inspect.isclass):
if m[0] not in _EXCLUDED_TRANSFORMS:
PIPELINES.register_module(m[1])
[docs]@PIPELINES.register_module(force=True)
class ToTensor(object):
"""Convert image or a sequence of images to tensor.
This module can not only convert a single image to tensor, but also a
sequence of images.
"""
def __init__(self) -> None:
self.transform = _transforms.ToTensor()
def __call__(self, imgs: Union[object, Sequence[object]]) -> torch.Tensor:
if isinstance(imgs, Sequence):
imgs = list(imgs)
for i, img in enumerate(imgs):
imgs[i] = self.transform(img)
else:
imgs = self.transform(imgs)
return imgs
[docs]@PIPELINES.register_module()
class SimMIMMaskGenerator(object):
"""Generate random block mask for each Image.
This module is used in SimMIM to generate masks.
Args:
input_size (int): Size of input image. Defaults to 192.
mask_patch_size (int): Size of each block mask. Defaults to 32.
model_patch_size (int): Patch size of each token. Defaults to 4.
mask_ratio (float): The mask ratio of image. Defaults to 0.6.
"""
def __init__(self,
input_size: int = 192,
mask_patch_size: int = 32,
model_patch_size: int = 4,
mask_ratio: float = 0.6) -> None:
self.input_size = input_size
self.mask_patch_size = mask_patch_size
self.model_patch_size = model_patch_size
self.mask_ratio = mask_ratio
assert self.input_size % self.mask_patch_size == 0
assert self.mask_patch_size % self.model_patch_size == 0
self.rand_size = self.input_size // self.mask_patch_size
self.scale = self.mask_patch_size // self.model_patch_size
self.token_count = self.rand_size**2
self.mask_count = int(np.ceil(self.token_count * self.mask_ratio))
def __call__(self, img: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
mask_idx = np.random.permutation(self.token_count)[:self.mask_count]
mask = np.zeros(self.token_count, dtype=int)
mask[mask_idx] = 1
mask = mask.reshape((self.rand_size, self.rand_size))
mask = mask.repeat(self.scale, axis=0).repeat(self.scale, axis=1)
mask = torch.from_numpy(mask) # H X W
return img, mask
[docs]@PIPELINES.register_module()
class BEiTMaskGenerator(object):
"""Generate mask for image.
This module is borrowed from
https://github.com/microsoft/unilm/tree/master/beit
Args:
input_size (int): The size of input image.
num_masking_patches (int): The number of patches to be masked.
min_num_patches (int): The minimum number of patches to be masked
in the process of generating mask. Defaults to 4.
max_num_patches (int, optional): The maximum number of patches to be
masked in the process of generating mask. Defaults to None.
min_aspect (float): The minimum aspect ratio of mask blocks. Defaults
to 0.3.
min_aspect (float, optional): The minimum aspect ratio of mask blocks.
Defaults to None.
"""
def __init__(self,
input_size: int,
num_masking_patches: int,
min_num_patches: int = 4,
max_num_patches: Optional[int] = None,
min_aspect: float = 0.3,
max_aspect: Optional[float] = None) -> None:
if not isinstance(input_size, tuple):
input_size = (input_size, ) * 2
self.height, self.width = input_size
self.num_patches = self.height * self.width
self.num_masking_patches = num_masking_patches
self.min_num_patches = min_num_patches
self.max_num_patches = num_masking_patches if max_num_patches is None \
else max_num_patches
max_aspect = max_aspect or 1 / min_aspect
self.log_aspect_ratio = (math.log(min_aspect), math.log(max_aspect))
def __repr__(self) -> None:
repr_str = 'Generator(%d, %d -> [%d ~ %d], max = %d, %.3f ~ %.3f)' % (
self.height, self.width, self.min_num_patches,
self.max_num_patches, self.num_masking_patches,
self.log_aspect_ratio[0], self.log_aspect_ratio[1])
return repr_str
def get_shape(self) -> Tuple[int, int]:
return self.height, self.width
def _mask(self, mask: np.ndarray, max_mask_patches: int) -> int:
delta = 0
for _ in range(10):
target_area = random.uniform(self.min_num_patches,
max_mask_patches)
aspect_ratio = math.exp(random.uniform(*self.log_aspect_ratio))
h = int(round(math.sqrt(target_area * aspect_ratio)))
w = int(round(math.sqrt(target_area / aspect_ratio)))
if w < self.width and h < self.height:
top = random.randint(0, self.height - h)
left = random.randint(0, self.width - w)
num_masked = mask[top:top + h, left:left + w].sum()
# Overlap
if 0 < h * w - num_masked <= max_mask_patches:
for i in range(top, top + h):
for j in range(left, left + w):
if mask[i, j] == 0:
mask[i, j] = 1
delta += 1
if delta > 0:
break
return delta
def __call__(
self, img: Tuple[torch.Tensor, torch.Tensor]
) -> Tuple[torch.Tensor, torch.Tensor, np.ndarray]:
mask = np.zeros(shape=self.get_shape(), dtype=int)
mask_count = 0
while mask_count != self.num_masking_patches:
max_mask_patches = self.num_masking_patches - mask_count
max_mask_patches = min(max_mask_patches, self.max_num_patches)
delta = self._mask(mask, max_mask_patches)
mask_count += delta
return img[0], img[1], mask
@PIPELINES.register_module()
class RandomResizedCropAndInterpolationWithTwoPic(object):
"""Crop the given PIL Image to random size and aspect ratio with random
interpolation.
This module is borrowed from
https://github.com/microsoft/unilm/tree/master/beit.
A crop of random size (default: of 0.08 to 1.0) of the original size and a
random aspect ratio (default: of 3/4 to 4/3) of the original aspect ratio
is made. This crop is finally resized to given size. This is popularly used
to train the Inception networks. This module first crops the image and
resizes the crop to two different sizes.
Args:
size (Union[tuple, int]): Expected output size of each edge of the
first image.
second_size (Union[tuple, int], optional): Expected output size of each
edge of the second image.
scale (tuple[float, float]): Range of size of the origin size cropped.
Defaults to (0.08, 1.0).
ratio (tuple[float, float]): Range of aspect ratio of the origin aspect
ratio cropped. Defaults to (3./4., 4./3.).
interpolation (str): The interpolation for the first image. Defaults
to ``bilinear``.
second_interpolation (str): The interpolation for the second image.
Defaults to ``lanczos``.
"""
interpolation_dict = {
'bicubic': Image.BICUBIC,
'lanczos': Image.LANCZOS,
'hamming': Image.HAMMING
}
def __init__(self,
size: Union[tuple, int],
second_size=None,
scale=(0.08, 1.0),
ratio=(3. / 4., 4. / 3.),
interpolation='bilinear',
second_interpolation='lanczos') -> None:
if isinstance(size, tuple):
self.size = size
else:
self.size = (size, size)
if second_size is not None:
if isinstance(second_size, tuple):
self.second_size = second_size
else:
self.second_size = (second_size, second_size)
else:
self.second_size = None
if (scale[0] > scale[1]) or (ratio[0] > ratio[1]):
warnings.warn('range should be of kind (min, max)')
if interpolation == 'random':
self.interpolation = (Image.BILINEAR, Image.BICUBIC)
else:
self.interpolation = self.interpolation_dict.get(
interpolation, Image.BILINEAR)
self.second_interpolation = self.interpolation_dict.get(
second_interpolation, Image.BILINEAR)
self.scale = scale
self.ratio = ratio
@staticmethod
def get_params(img: np.ndarray, scale: tuple,
ratio: tuple) -> Sequence[int]:
"""Get parameters for ``crop`` for a random sized crop.
Args:
img (np.ndarray): Image to be cropped.
scale (tuple): range of size of the origin size cropped
ratio (tuple): range of aspect ratio of the origin aspect
ratio cropped
Returns:
tuple: params (i, j, h, w) to be passed to ``crop`` for a random
sized crop.
"""
area = img.size[0] * img.size[1]
for _ in range(10):
target_area = random.uniform(*scale) * area
log_ratio = (math.log(ratio[0]), math.log(ratio[1]))
aspect_ratio = math.exp(random.uniform(*log_ratio))
w = int(round(math.sqrt(target_area * aspect_ratio)))
h = int(round(math.sqrt(target_area / aspect_ratio)))
if w <= img.size[0] and h <= img.size[1]:
i = random.randint(0, img.size[1] - h)
j = random.randint(0, img.size[0] - w)
return i, j, h, w
# Fallback to central crop
in_ratio = img.size[0] / img.size[1]
if in_ratio < min(ratio):
w = img.size[0]
h = int(round(w / min(ratio)))
elif in_ratio > max(ratio):
h = img.size[1]
w = int(round(h * max(ratio)))
else: # whole image
w = img.size[0]
h = img.size[1]
i = (img.size[1] - h) // 2
j = (img.size[0] - w) // 2
return i, j, h, w
def __call__(
self, img: np.ndarray
) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
i, j, h, w = self.get_params(img, self.scale, self.ratio)
if isinstance(self.interpolation, (tuple, list)):
interpolation = random.choice(self.interpolation)
else:
interpolation = self.interpolation
if self.second_size is None:
return F.resized_crop(img, i, j, h, w, self.size, interpolation)
else:
return F.resized_crop(img, i, j, h, w, self.size,
interpolation), F.resized_crop(
img, i, j, h, w, self.second_size,
self.second_interpolation)
[docs]@PIPELINES.register_module()
class RandomAug(object):
"""RandAugment data augmentation method based on
`"RandAugment: Practical automated data augmentation
with a reduced search space"
<https://arxiv.org/abs/1909.13719>`_.
This code is borrowed from <https://github.com/pengzhiliang/MAE-pytorch>
"""
def __init__(self,
input_size=None,
color_jitter=None,
auto_augment=None,
interpolation=None,
re_prob=None,
re_mode=None,
re_count=None,
mean=None,
std=None):
self.trans = create_transform(
input_size=input_size,
is_training=True,
color_jitter=color_jitter,
auto_augment=auto_augment,
interpolation=interpolation,
re_prob=re_prob,
re_mode=re_mode,
re_count=re_count,
mean=mean,
std=std,
)
def __call__(self, img):
return self.trans(img)
def __repr__(self) -> str:
repr_str = self.__class__.__name__
return repr_str
[docs]@PIPELINES.register_module()
class RandomAppliedTrans(object):
"""Randomly applied transformations.
Args:
transforms (list[dict]): List of transformations in dictionaries.
p (float, optional): Probability. Defaults to 0.5.
"""
def __init__(self, transforms, p=0.5):
t = [build_from_cfg(t, PIPELINES) for t in transforms]
self.trans = _transforms.RandomApply(t, p=p)
self.prob = p
def __call__(self, img):
return self.trans(img)
def __repr__(self):
repr_str = self.__class__.__name__
repr_str += f'prob = {self.prob}'
return repr_str
# custom transforms
[docs]@PIPELINES.register_module()
class Lighting(object):
"""Lighting noise(AlexNet - style PCA - based noise).
Args:
alphastd (float, optional): The parameter for Lighting.
Defaults to 0.1.
"""
_IMAGENET_PCA = {
'eigval':
torch.Tensor([0.2175, 0.0188, 0.0045]),
'eigvec':
torch.Tensor([
[-0.5675, 0.7192, 0.4009],
[-0.5808, -0.0045, -0.8140],
[-0.5836, -0.6948, 0.4203],
])
}
def __init__(self, alphastd=0.1):
self.alphastd = alphastd
self.eigval = self._IMAGENET_PCA['eigval']
self.eigvec = self._IMAGENET_PCA['eigvec']
def __call__(self, img):
assert isinstance(img, torch.Tensor), \
f'Expect torch.Tensor, got {type(img)}'
if self.alphastd == 0:
return img
alpha = img.new().resize_(3).normal_(0, self.alphastd)
rgb = self.eigvec.type_as(img).clone()\
.mul(alpha.view(1, 3).expand(3, 3))\
.mul(self.eigval.view(1, 3).expand(3, 3))\
.sum(1).squeeze()
return img.add(rgb.view(3, 1, 1).expand_as(img))
def __repr__(self):
repr_str = self.__class__.__name__
repr_str += f'alphastd = {self.alphastd}'
return repr_str
[docs]@PIPELINES.register_module()
class GaussianBlur(object):
"""GaussianBlur augmentation refers to `SimCLR.
<https://arxiv.org/abs/2002.05709>`_.
Args:
sigma_min (float): The minimum parameter of Gaussian kernel std.
sigma_max (float): The maximum parameter of Gaussian kernel std.
p (float, optional): Probability. Defaults to 0.5.
"""
def __init__(self, sigma_min, sigma_max, p=0.5):
assert 0 <= p <= 1.0, \
f'The prob should be in range [0,1], got {p} instead.'
self.sigma_min = sigma_min
self.sigma_max = sigma_max
self.prob = p
def __call__(self, img):
if np.random.rand() > self.prob:
return img
sigma = np.random.uniform(self.sigma_min, self.sigma_max)
img = img.filter(ImageFilter.GaussianBlur(radius=sigma))
return img
def __repr__(self):
repr_str = self.__class__.__name__
repr_str += f'sigma_min = {self.sigma_min}, '
repr_str += f'sigma_max = {self.sigma_max}, '
repr_str += f'prob = {self.prob}'
return repr_str
[docs]@PIPELINES.register_module()
class Solarization(object):
"""Solarization augmentation refers to `BYOL.
<https://arxiv.org/abs/2006.07733>`_.
Args:
threshold (float, optional): The solarization threshold.
Defaults to 128.
p (float, optional): Probability. Defaults to 0.5.
"""
def __init__(self, threshold=128, p=0.5):
assert 0 <= p <= 1.0, \
f'The prob should be in range [0, 1], got {p} instead.'
self.threshold = threshold
self.prob = p
def __call__(self, img):
if np.random.rand() > self.prob:
return img
img = np.array(img)
img = np.where(img < self.threshold, img, 255 - img)
return Image.fromarray(img.astype(np.uint8))
def __repr__(self):
repr_str = self.__class__.__name__
repr_str += f'threshold = {self.threshold}, '
repr_str += f'prob = {self.prob}'
return repr_str
[docs]@PIPELINES.register_module()
class MaskFeatMaskGenerator(object):
"""Generate random block mask for each image.
This module is borrowed from
https://github.com/facebookresearch/SlowFast/blob/main/slowfast/datasets/transform.py
Args:
mask_window_size (int): Size of input image. Defaults to 14.
mask_ratio (float): The mask ratio of image. Defaults to 0.4.
min_num_patches (int): Minimum number of patches that require masking.
Defaults to 15.
max_num_patches (int, optional): Maximum number of patches that
require masking. Defaults to None.
min_aspect (int): Minimum aspect of patches. Defaults to 0.3.
max_aspect (float, optional): Maximum aspect of patches.
Defaults to None.
"""
def __init__(
self,
mask_window_size: int = 14,
mask_ratio: float = 0.4,
min_num_patches: int = 15,
max_num_patches: Optional[int] = None,
min_aspect: float = 0.3,
max_aspect: Optional[float] = None,
) -> None:
self.height, self.width = mask_window_size, mask_window_size
self.num_patches = self.height * self.width
self.num_masking_patches = int(mask_window_size**2 * mask_ratio)
self.min_num_patches = min_num_patches
self.max_num_patches = (
self.num_masking_patches
if max_num_patches is None else max_num_patches)
max_aspect = max_aspect or 1 / min_aspect
self.log_aspect_ratio = (math.log(min_aspect), math.log(max_aspect))
def __repr__(self) -> str:
repr_str = self.__class__.__name__
repr_str += f'(height={self.height}, '
repr_str += f'width={self.width}, '
repr_str += f'num_patches={self.num_patches}, '
repr_str += f'num_masking_patches={self.num_masking_patches}, '
repr_str += f'min_num_patches={self.min_num_patches}, '
repr_str += f'max_num_patches={self.max_num_patches}, '
repr_str += f'log_aspect_ratio={self.log_aspect_ratio})'
return repr_str
def get_shape(self) -> Tuple[int, int]:
return self.height, self.width
def _random_masking(self, mask: np.array, max_mask_patches: int) -> int:
"""Generate random block masks for each image up to 10 times.
Args:
mask (np.array): Initial mask of shape (mask_window_size,
mask_window_size).
max_mask_patches (int): Maximum number of masked patches required.
Returns:
int: Number of masking patches.
"""
delta = 0
for _ in range(10):
target_area = random.uniform(self.min_num_patches,
max_mask_patches)
aspect_ratio = math.exp(random.uniform(*self.log_aspect_ratio))
h = int(round(math.sqrt(target_area * aspect_ratio)))
w = int(round(math.sqrt(target_area / aspect_ratio)))
if w < self.width and h < self.height:
top = random.randint(0, self.height - h)
left = random.randint(0, self.width - w)
num_masked = mask[top:top + h, left:left + w].sum()
# Overlap
if 0 < h * w - num_masked <= max_mask_patches:
for i in range(top, top + h):
for j in range(left, left + w):
if mask[i, j] == 0:
mask[i, j] = 1
delta += 1
if delta > 0:
break
return delta
def __call__(self, img: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
"""Generate random block mask for each image.
Args:
img (torch.Tensor): Input image of shape (C, H, W).
Returns:
Tuple[torch.Tensor, torch.Tensor]: Input image and mask.
"""
mask = np.zeros(shape=self.get_shape(), dtype=int)
mask_count = 0
while mask_count < self.num_masking_patches:
max_mask_patches = self.num_masking_patches - mask_count
max_mask_patches = min(max_mask_patches, self.max_num_patches)
delta = self._random_masking(mask, max_mask_patches)
if delta == 0:
break
else:
mask_count += delta
return img, torch.Tensor(mask).bool()