You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
251 lines
8.4 KiB
251 lines
8.4 KiB
import torch
|
|
from torchvision import transforms
|
|
import torchvision.transforms.functional as F
|
|
from PIL import Image
|
|
import warnings
|
|
import math
|
|
import random
|
|
import numpy as np
|
|
|
|
from .constants import DEFAULT_CROP_PCT, IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
|
|
from .random_erasing import RandomErasing
|
|
from .auto_augment import AutoAugment, auto_augment_policy
|
|
|
|
|
|
class ToNumpy:
|
|
|
|
def __call__(self, pil_img):
|
|
np_img = np.array(pil_img, dtype=np.uint8)
|
|
if np_img.ndim < 3:
|
|
np_img = np.expand_dims(np_img, axis=-1)
|
|
np_img = np.rollaxis(np_img, 2) # HWC to CHW
|
|
return np_img
|
|
|
|
|
|
class ToTensor:
|
|
|
|
def __init__(self, dtype=torch.float32):
|
|
self.dtype = dtype
|
|
|
|
def __call__(self, pil_img):
|
|
np_img = np.array(pil_img, dtype=np.uint8)
|
|
if np_img.ndim < 3:
|
|
np_img = np.expand_dims(np_img, axis=-1)
|
|
np_img = np.rollaxis(np_img, 2) # HWC to CHW
|
|
return torch.from_numpy(np_img).to(dtype=self.dtype)
|
|
|
|
|
|
_pil_interpolation_to_str = {
|
|
Image.NEAREST: 'PIL.Image.NEAREST',
|
|
Image.BILINEAR: 'PIL.Image.BILINEAR',
|
|
Image.BICUBIC: 'PIL.Image.BICUBIC',
|
|
Image.LANCZOS: 'PIL.Image.LANCZOS',
|
|
Image.HAMMING: 'PIL.Image.HAMMING',
|
|
Image.BOX: 'PIL.Image.BOX',
|
|
}
|
|
|
|
|
|
def _pil_interp(method):
|
|
if method == 'bicubic':
|
|
return Image.BICUBIC
|
|
elif method == 'lanczos':
|
|
return Image.LANCZOS
|
|
elif method == 'hamming':
|
|
return Image.HAMMING
|
|
else:
|
|
# default bilinear, do we want to allow nearest?
|
|
return Image.BILINEAR
|
|
|
|
|
|
_RANDOM_INTERPOLATION = (Image.BILINEAR, Image.BICUBIC)
|
|
|
|
|
|
class RandomResizedCropAndInterpolation:
|
|
"""Crop the given PIL Image to random size and aspect ratio with random interpolation.
|
|
|
|
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.
|
|
|
|
Args:
|
|
size: expected output size of each edge
|
|
scale: range of size of the origin size cropped
|
|
ratio: range of aspect ratio of the origin aspect ratio cropped
|
|
interpolation: Default: PIL.Image.BILINEAR
|
|
"""
|
|
|
|
def __init__(self, size, scale=(0.08, 1.0), ratio=(3. / 4., 4. / 3.),
|
|
interpolation='bilinear'):
|
|
if isinstance(size, tuple):
|
|
self.size = size
|
|
else:
|
|
self.size = (size, size)
|
|
if (scale[0] > scale[1]) or (ratio[0] > ratio[1]):
|
|
warnings.warn("range should be of kind (min, max)")
|
|
|
|
if interpolation == 'random':
|
|
self.interpolation = _RANDOM_INTERPOLATION
|
|
else:
|
|
self.interpolation = _pil_interp(interpolation)
|
|
self.scale = scale
|
|
self.ratio = ratio
|
|
|
|
@staticmethod
|
|
def get_params(img, scale, ratio):
|
|
"""Get parameters for ``crop`` for a random sized crop.
|
|
|
|
Args:
|
|
img (PIL Image): 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 attempt 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):
|
|
"""
|
|
Args:
|
|
img (PIL Image): Image to be cropped and resized.
|
|
|
|
Returns:
|
|
PIL Image: Randomly cropped and resized image.
|
|
"""
|
|
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
|
|
return F.resized_crop(img, i, j, h, w, self.size, interpolation)
|
|
|
|
def __repr__(self):
|
|
if isinstance(self.interpolation, (tuple, list)):
|
|
interpolate_str = ' '.join([_pil_interpolation_to_str[x] for x in self.interpolation])
|
|
else:
|
|
interpolate_str = _pil_interpolation_to_str[self.interpolation]
|
|
format_string = self.__class__.__name__ + '(size={0}'.format(self.size)
|
|
format_string += ', scale={0}'.format(tuple(round(s, 4) for s in self.scale))
|
|
format_string += ', ratio={0}'.format(tuple(round(r, 4) for r in self.ratio))
|
|
format_string += ', interpolation={0})'.format(interpolate_str)
|
|
return format_string
|
|
|
|
|
|
def transforms_imagenet_train(
|
|
img_size=224,
|
|
scale=(0.08, 1.0),
|
|
color_jitter=0.4,
|
|
auto_augment=None,
|
|
interpolation='random',
|
|
random_erasing=0.4,
|
|
random_erasing_mode='const',
|
|
use_prefetcher=False,
|
|
mean=IMAGENET_DEFAULT_MEAN,
|
|
std=IMAGENET_DEFAULT_STD
|
|
):
|
|
tfl = [
|
|
RandomResizedCropAndInterpolation(
|
|
img_size, scale=scale, interpolation=interpolation),
|
|
transforms.RandomHorizontalFlip()
|
|
]
|
|
if auto_augment:
|
|
aa_params = dict(
|
|
translate_const=img_size[-1] // 2 - 1,
|
|
img_mean=tuple([min(255, round(255 * x)) for x in mean]),
|
|
)
|
|
if interpolation and interpolation != 'random':
|
|
aa_params['interpolation'] = _pil_interp(interpolation)
|
|
aa_policy = auto_augment_policy(auto_augment, aa_params)
|
|
tfl += [AutoAugment(aa_policy)]
|
|
else:
|
|
# color jitter is enabled when not using AA
|
|
if isinstance(color_jitter, (list, tuple)):
|
|
# color jitter should be a 3-tuple/list if spec brightness/contrast/saturation
|
|
# or 4 if also augmenting hue
|
|
assert len(color_jitter) in (3, 4)
|
|
else:
|
|
# if it's a scalar, duplicate for brightness, contrast, and saturation, no hue
|
|
color_jitter = (float(color_jitter),) * 3
|
|
tfl += [transforms.ColorJitter(*color_jitter)]
|
|
|
|
if use_prefetcher:
|
|
# prefetcher and collate will handle tensor conversion and norm
|
|
tfl += [ToNumpy()]
|
|
else:
|
|
tfl += [
|
|
transforms.ToTensor(),
|
|
transforms.Normalize(
|
|
mean=torch.tensor(mean),
|
|
std=torch.tensor(std))
|
|
]
|
|
if random_erasing > 0.:
|
|
tfl.append(RandomErasing(random_erasing, mode=random_erasing_mode, device='cpu'))
|
|
return transforms.Compose(tfl)
|
|
|
|
|
|
def transforms_imagenet_eval(
|
|
img_size=224,
|
|
crop_pct=None,
|
|
interpolation='bilinear',
|
|
use_prefetcher=False,
|
|
mean=IMAGENET_DEFAULT_MEAN,
|
|
std=IMAGENET_DEFAULT_STD):
|
|
crop_pct = crop_pct or DEFAULT_CROP_PCT
|
|
|
|
if isinstance(img_size, tuple):
|
|
assert len(img_size) == 2
|
|
if img_size[-1] == img_size[-2]:
|
|
# fall-back to older behaviour so Resize scales to shortest edge if target is square
|
|
scale_size = int(math.floor(img_size[0] / crop_pct))
|
|
else:
|
|
scale_size = tuple([int(x / crop_pct) for x in img_size])
|
|
else:
|
|
scale_size = int(math.floor(img_size / crop_pct))
|
|
|
|
tfl = [
|
|
transforms.Resize(scale_size, _pil_interp(interpolation)),
|
|
transforms.CenterCrop(img_size),
|
|
]
|
|
if use_prefetcher:
|
|
# prefetcher and collate will handle tensor conversion and norm
|
|
tfl += [ToNumpy()]
|
|
else:
|
|
tfl += [
|
|
transforms.ToTensor(),
|
|
transforms.Normalize(
|
|
mean=torch.tensor(mean),
|
|
std=torch.tensor(std))
|
|
]
|
|
|
|
return transforms.Compose(tfl)
|