|
|
|
|
@ -1,3 +1,9 @@
|
|
|
|
|
import math
|
|
|
|
|
import numbers
|
|
|
|
|
import random
|
|
|
|
|
import warnings
|
|
|
|
|
from typing import List, Sequence
|
|
|
|
|
|
|
|
|
|
import torch
|
|
|
|
|
import torchvision.transforms.functional as F
|
|
|
|
|
try:
|
|
|
|
|
@ -6,9 +12,6 @@ try:
|
|
|
|
|
except ImportError:
|
|
|
|
|
has_interpolation_mode = False
|
|
|
|
|
from PIL import Image
|
|
|
|
|
import warnings
|
|
|
|
|
import math
|
|
|
|
|
import random
|
|
|
|
|
import numpy as np
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@ -96,6 +99,19 @@ def interp_mode_to_str(mode):
|
|
|
|
|
_RANDOM_INTERPOLATION = (str_to_interp_mode('bilinear'), str_to_interp_mode('bicubic'))
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _setup_size(size, error_msg):
|
|
|
|
|
if isinstance(size, numbers.Number):
|
|
|
|
|
return int(size), int(size)
|
|
|
|
|
|
|
|
|
|
if isinstance(size, Sequence) and len(size) == 1:
|
|
|
|
|
return size[0], size[0]
|
|
|
|
|
|
|
|
|
|
if len(size) != 2:
|
|
|
|
|
raise ValueError(error_msg)
|
|
|
|
|
|
|
|
|
|
return size
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class RandomResizedCropAndInterpolation:
|
|
|
|
|
"""Crop the given PIL Image to random size and aspect ratio with random interpolation.
|
|
|
|
|
|
|
|
|
|
@ -195,3 +211,132 @@ class RandomResizedCropAndInterpolation:
|
|
|
|
|
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 center_crop_or_pad(img: torch.Tensor, output_size: List[int], fill=0) -> torch.Tensor:
|
|
|
|
|
"""Center crops and/or pads the given image.
|
|
|
|
|
If the image is torch Tensor, it is expected
|
|
|
|
|
to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions.
|
|
|
|
|
If image size is smaller than output size along any edge, image is padded with 0 and then center cropped.
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
img (PIL Image or Tensor): Image to be cropped.
|
|
|
|
|
output_size (sequence or int): (height, width) of the crop box. If int or sequence with single int,
|
|
|
|
|
it is used for both directions.
|
|
|
|
|
fill (int, Tuple[int]): Padding color
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
PIL Image or Tensor: Cropped image.
|
|
|
|
|
"""
|
|
|
|
|
if isinstance(output_size, numbers.Number):
|
|
|
|
|
output_size = (int(output_size), int(output_size))
|
|
|
|
|
elif isinstance(output_size, (tuple, list)) and len(output_size) == 1:
|
|
|
|
|
output_size = (output_size[0], output_size[0])
|
|
|
|
|
|
|
|
|
|
_, image_height, image_width = F.get_dimensions(img)
|
|
|
|
|
crop_height, crop_width = output_size
|
|
|
|
|
|
|
|
|
|
if crop_width > image_width or crop_height > image_height:
|
|
|
|
|
padding_ltrb = [
|
|
|
|
|
(crop_width - image_width) // 2 if crop_width > image_width else 0,
|
|
|
|
|
(crop_height - image_height) // 2 if crop_height > image_height else 0,
|
|
|
|
|
(crop_width - image_width + 1) // 2 if crop_width > image_width else 0,
|
|
|
|
|
(crop_height - image_height + 1) // 2 if crop_height > image_height else 0,
|
|
|
|
|
]
|
|
|
|
|
img = F.pad(img, padding_ltrb, fill=fill)
|
|
|
|
|
_, image_height, image_width = F.get_dimensions(img)
|
|
|
|
|
if crop_width == image_width and crop_height == image_height:
|
|
|
|
|
return img
|
|
|
|
|
|
|
|
|
|
crop_top = int(round((image_height - crop_height) / 2.0))
|
|
|
|
|
crop_left = int(round((image_width - crop_width) / 2.0))
|
|
|
|
|
return F.crop(img, crop_top, crop_left, crop_height, crop_width)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class CenterCropOrPad(torch.nn.Module):
|
|
|
|
|
"""Crops the given image at the center.
|
|
|
|
|
If the image is torch Tensor, it is expected
|
|
|
|
|
to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions.
|
|
|
|
|
If image size is smaller than output size along any edge, image is padded with 0 and then center cropped.
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
size (sequence or int): Desired output size of the crop. If size is an
|
|
|
|
|
int instead of sequence like (h, w), a square crop (size, size) is
|
|
|
|
|
made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]).
|
|
|
|
|
"""
|
|
|
|
|
|
|
|
|
|
def __init__(self, size, fill=0):
|
|
|
|
|
super().__init__()
|
|
|
|
|
self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.")
|
|
|
|
|
self.fill = fill
|
|
|
|
|
|
|
|
|
|
def forward(self, img):
|
|
|
|
|
"""
|
|
|
|
|
Args:
|
|
|
|
|
img (PIL Image or Tensor): Image to be cropped.
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
PIL Image or Tensor: Cropped image.
|
|
|
|
|
"""
|
|
|
|
|
return center_crop_or_pad(img, self.size, fill=self.fill)
|
|
|
|
|
|
|
|
|
|
def __repr__(self) -> str:
|
|
|
|
|
return f"{self.__class__.__name__}(size={self.size})"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class ResizeKeepRatio:
|
|
|
|
|
""" Resize and Keep Ratio
|
|
|
|
|
"""
|
|
|
|
|
|
|
|
|
|
def __init__(
|
|
|
|
|
self,
|
|
|
|
|
size,
|
|
|
|
|
longest=0.,
|
|
|
|
|
interpolation='bilinear',
|
|
|
|
|
fill=0,
|
|
|
|
|
):
|
|
|
|
|
if isinstance(size, (list, tuple)):
|
|
|
|
|
self.size = tuple(size)
|
|
|
|
|
else:
|
|
|
|
|
self.size = (size, size)
|
|
|
|
|
self.interpolation = str_to_interp_mode(interpolation)
|
|
|
|
|
self.longest = float(longest)
|
|
|
|
|
self.fill = fill
|
|
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
|
def get_params(img, target_size, longest):
|
|
|
|
|
"""Get parameters
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
img (PIL Image): Image to be cropped.
|
|
|
|
|
target_size (Tuple[int, int]): Size of output
|
|
|
|
|
Returns:
|
|
|
|
|
tuple: params (h, w) and (l, r, t, b) to be passed to ``resize`` and ``pad`` respectively
|
|
|
|
|
"""
|
|
|
|
|
source_size = img.size[::-1] # h, w
|
|
|
|
|
h, w = source_size
|
|
|
|
|
target_h, target_w = target_size
|
|
|
|
|
ratio_h = h / target_h
|
|
|
|
|
ratio_w = w / target_w
|
|
|
|
|
ratio = max(ratio_h, ratio_w) * longest + min(ratio_h, ratio_w) * (1. - longest)
|
|
|
|
|
size = [round(x / ratio) for x in source_size]
|
|
|
|
|
return size
|
|
|
|
|
|
|
|
|
|
def __call__(self, img):
|
|
|
|
|
"""
|
|
|
|
|
Args:
|
|
|
|
|
img (PIL Image): Image to be cropped and resized.
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
PIL Image: Resized, padded to at least target size, possibly cropped to exactly target size
|
|
|
|
|
"""
|
|
|
|
|
size = self.get_params(img, self.size, self.longest)
|
|
|
|
|
img = F.resize(img, size, self.interpolation)
|
|
|
|
|
return img
|
|
|
|
|
|
|
|
|
|
def __repr__(self):
|
|
|
|
|
interpolate_str = interp_mode_to_str(self.interpolation)
|
|
|
|
|
format_string = self.__class__.__name__ + '(size={0}'.format(self.size)
|
|
|
|
|
format_string += f', interpolation={interpolate_str})'
|
|
|
|
|
format_string += f', longest={self.longest:.3f})'
|
|
|
|
|
return format_string
|
|
|
|
|
|
Ross Wightman
1 year ago