fastai--fastai
wehub-resource-sync
1269 行
56 KiB
Python
1269 行
56 KiB
Python
"""Transforms to apply data augmentation in Computer Vision
|
|
|
|
Docs: https://docs.fast.ai/vision.augment.html.md"""
|
|
|
|
# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/09_vision.augment.ipynb.
|
|
|
|
# %% auto #0
|
|
__all__ = ['TensorTypes', 'RandTransform', 'FlipItem', 'DihedralItem', 'CropPad', 'RandomCrop', 'OldRandomCrop', 'Resize',
|
|
'RandomResizedCrop', 'RatioResize', 'affine_grid', 'AffineCoordTfm', 'RandomResizedCropGPU', 'mask_tensor',
|
|
'affine_mat', 'flip_mat', 'Flip', 'DeterministicDraw', 'DeterministicFlip', 'dihedral_mat', 'Dihedral',
|
|
'DeterministicDihedral', 'rotate_mat', 'Rotate', 'zoom_mat', 'Zoom', 'find_coeffs', 'apply_perspective',
|
|
'Warp', 'SpaceTfm', 'LightingTfm', 'Brightness', 'Contrast', 'grayscale', 'Saturation', 'rgb2hsv', 'hsv2rgb',
|
|
'HSVTfm', 'Hue', 'cutout_gaussian', 'norm_apply_denorm', 'RandomErasing', 'setup_aug_tfms', 'aug_transforms',
|
|
'PadMode', 'ResizeMethod']
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #bcd15f9f
|
|
from ..data.all import *
|
|
from .core import *
|
|
from .data import *
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #4f678d91
|
|
from torch import stack, zeros_like as t0, ones_like as t1
|
|
from torch.distributions.bernoulli import Bernoulli
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #d5a5551d
|
|
class RandTransform(DisplayedTransform):
|
|
"A transform that before_call its state at each `__call__`"
|
|
do,nm,supports,split_idx = True,None,[],0
|
|
def __init__(self,
|
|
p:float=1., # Probability of applying Transform
|
|
nm:str=None,
|
|
before_call:Callable=None, # Optional batchwise preprocessing function
|
|
**kwargs
|
|
):
|
|
store_attr('p')
|
|
super().__init__(**kwargs)
|
|
self.before_call = ifnone(before_call,self.before_call)
|
|
|
|
def before_call(self,
|
|
b,
|
|
split_idx:int, # Index of the train/valid dataset
|
|
):
|
|
"This function can be overridden. Set `self.do` based on `self.p`"
|
|
self.do = self.p==1. or random.random() < self.p
|
|
|
|
def __call__(self,
|
|
b,
|
|
split_idx:int=None, # Index of the train/valid dataset
|
|
**kwargs
|
|
):
|
|
self.before_call(b, split_idx=split_idx)
|
|
return super().__call__(b, split_idx=split_idx, **kwargs) if self.do else b
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #2d71db2a
|
|
def _neg_axis(x, axis):
|
|
x[...,axis] = -x[...,axis]
|
|
return x
|
|
|
|
TensorTypes = (TensorImage,TensorMask,TensorPoint,TensorBBox)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #a833f80e
|
|
@patch
|
|
def flip_lr(x:Image.Image): return x.transpose(Image.FLIP_LEFT_RIGHT)
|
|
@patch
|
|
def flip_lr(x:TensorImageBase): return x.flip(-1)
|
|
@patch
|
|
def flip_lr(x:TensorPoint): return TensorPoint(_neg_axis(x.clone(), 0))
|
|
@patch
|
|
def flip_lr(x:TensorBBox): return TensorBBox(TensorPoint(x.view(-1,2)).flip_lr().view(-1,4))
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #b7bd957b
|
|
class FlipItem(RandTransform):
|
|
"Randomly flip with probability `p`"
|
|
def __init__(self, p:float=0.5): super().__init__(p=p)
|
|
def encodes(self, x:(Image.Image,*TensorTypes)): return x.flip_lr()
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #625b9e24
|
|
@patch
|
|
def dihedral(x:PILImage,
|
|
k:int, # Dihedral transformation to apply
|
|
):
|
|
return x if k==0 else x.transpose(k-1)
|
|
@patch
|
|
def dihedral(x:TensorImage,
|
|
k:int, # Dihedral transformation to apply
|
|
):
|
|
if k in [1,3,4,7]: x = x.flip(-1)
|
|
if k in [2,4,5,7]: x = x.flip(-2)
|
|
if k in [3,5,6,7]: x = x.transpose(-1,-2)
|
|
return x
|
|
@patch
|
|
def dihedral(x:TensorPoint,
|
|
k:int, # Dihedral transformation to apply
|
|
):
|
|
if k in [1,3,4,7]: x = _neg_axis(x, 0)
|
|
if k in [2,4,5,7]: x = _neg_axis(x, 1)
|
|
if k in [3,5,6,7]: x = x.flip(1)
|
|
return x
|
|
@patch
|
|
def dihedral(x:TensorBBox,
|
|
k:int, #Dihedral transformation to apply
|
|
):
|
|
pnts = TensorPoint(x.view(-1,2)).dihedral(k).view(-1,2,2)
|
|
tl,br = pnts.min(dim=1)[0],pnts.max(dim=1)[0]
|
|
return TensorBBox(torch.cat([tl, br], dim=1), img_size=x.img_size)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #52182f5b
|
|
class DihedralItem(RandTransform):
|
|
"Randomly flip with probability `p`"
|
|
def before_call(self, b, split_idx):
|
|
super().before_call(b, split_idx)
|
|
self.k = random.randint(0,7)
|
|
|
|
def encodes(self, x:(Image.Image,*TensorTypes)): return x.dihedral(self.k)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #b55a5ac8
|
|
from torchvision.transforms.functional import pad as tvpad
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #01a824b0
|
|
mk_class('PadMode', **{o:o.lower() for o in ['Zeros', 'Border', 'Reflection']},
|
|
doc="All possible padding mode as attributes to get tab-completion and typo-proofing")
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #80f293fe
|
|
_all_ = ['PadMode']
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #1c2876b7
|
|
_pad_modes = {'zeros': 'constant', 'border': 'edge', 'reflection': 'reflect'}
|
|
|
|
@patch
|
|
def _do_crop_pad(x:Image.Image, sz, tl, orig_sz,
|
|
pad_mode=PadMode.Zeros, resize_mode=BILINEAR, resize_to=None):
|
|
if any(tl.ge(0)) or any(tl.add(sz).le(orig_sz)):
|
|
# At least one dim is inside the image, so needs to be cropped
|
|
c = tl.max(0)
|
|
x = x.crop((*c, *tl.add(sz).min(orig_sz)))
|
|
if any(tl.lt(0)) or any(tl.add(sz).ge(orig_sz)):
|
|
# At least one dim is outside the image, so needs to be padded
|
|
p = (-tl).max(0)
|
|
f = (sz-orig_sz).add(tl).max(0)
|
|
x = tvpad(x, (*p, *f), padding_mode=_pad_modes[pad_mode])
|
|
if resize_to is not None: x = x.resize(resize_to, resize_mode)
|
|
return x
|
|
|
|
@patch
|
|
def _do_crop_pad(x:TensorPoint, sz, tl, orig_sz, pad_mode=PadMode.Zeros, resize_to=None, **kwargs):
|
|
#assert pad_mode==PadMode.Zeros,"Only zero padding is supported for `TensorPoint` and `TensorBBox`"
|
|
orig_sz,sz,tl = map(FloatTensor, (orig_sz,sz,tl))
|
|
return TensorPoint((x+1)*orig_sz/sz - tl*2/sz - 1, sz=sz if resize_to is None else resize_to)
|
|
|
|
@patch
|
|
def _do_crop_pad(x:TensorBBox, sz, tl, orig_sz, pad_mode=PadMode.Zeros, resize_to=None, **kwargs):
|
|
bbox = TensorPoint._do_crop_pad(x.view(-1,2), sz, tl, orig_sz, pad_mode, resize_to).view(-1,4)
|
|
return TensorBBox(bbox, img_size=x.img_size)
|
|
|
|
@patch
|
|
def crop_pad(x:TensorBBox|TensorPoint|Image.Image,
|
|
sz:int|tuple, # Crop/pad size of input, duplicated if one value is specified
|
|
tl:tuple=None, # Optional top-left coordinate of the crop/pad, if `None` center crop
|
|
orig_sz:tuple=None, # Original size of input
|
|
pad_mode:PadMode=PadMode.Zeros, # Fastai padding mode
|
|
resize_mode=BILINEAR, # Pillow `Image` resize mode
|
|
resize_to:tuple=None # Optional post crop/pad resize of input
|
|
):
|
|
if isinstance(sz,int): sz = (sz,sz)
|
|
orig_sz = fastuple(_get_sz(x) if orig_sz is None else orig_sz)
|
|
sz,tl = fastuple(sz),fastuple(((_get_sz(x)-sz)//2) if tl is None else tl)
|
|
return x._do_crop_pad(sz, tl, orig_sz=orig_sz, pad_mode=pad_mode, resize_mode=resize_mode, resize_to=resize_to)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #f29647dd
|
|
def _process_sz(size):
|
|
if isinstance(size,int): size=(size,size)
|
|
return fastuple(size[1],size[0])
|
|
|
|
def _get_sz(x):
|
|
if isinstance(x, tuple): x = x[0]
|
|
if not isinstance(x, Tensor): return fastuple(x.size)
|
|
return fastuple(getattr(x, 'img_size', getattr(x, 'sz', (x.shape[-1], x.shape[-2]))))
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #c0553fdc
|
|
@delegates()
|
|
class CropPad(DisplayedTransform):
|
|
"Center crop or pad an image to `size`"
|
|
order = 0
|
|
def __init__(self,
|
|
size:int|tuple, # Size to crop or pad to, duplicated if one value is specified
|
|
pad_mode:PadMode=PadMode.Zeros, # A `PadMode`
|
|
**kwargs
|
|
):
|
|
size = _process_sz(size)
|
|
store_attr()
|
|
super().__init__(**kwargs)
|
|
|
|
def encodes(self, x:Image.Image|TensorBBox|TensorPoint):
|
|
orig_sz = _get_sz(x)
|
|
tl = (orig_sz-self.size)//2
|
|
return x.crop_pad(self.size, tl, orig_sz=orig_sz, pad_mode=self.pad_mode)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #74b79240
|
|
@delegates()
|
|
class RandomCrop(RandTransform):
|
|
"Randomly crop an image to `size`"
|
|
split_idx,order = None,1
|
|
def __init__(self,
|
|
size:int|tuple, # Size to crop to, duplicated if one value is specified
|
|
**kwargs
|
|
):
|
|
size = _process_sz(size)
|
|
store_attr()
|
|
super().__init__(**kwargs)
|
|
|
|
def before_call(self,
|
|
b,
|
|
split_idx:int # Index of the train/valid dataset
|
|
):
|
|
"Randomly positioning crop if train dataset else center crop"
|
|
self.orig_sz = _get_sz(b)
|
|
if split_idx: self.tl = (self.orig_sz-self.size)//2
|
|
else:
|
|
wd = self.orig_sz[0] - self.size[0]
|
|
hd = self.orig_sz[1] - self.size[1]
|
|
w_rand = (wd, -1) if wd < 0 else (0, wd)
|
|
h_rand = (hd, -1) if hd < 0 else (0, hd)
|
|
self.tl = fastuple(random.randint(*w_rand), random.randint(*h_rand))
|
|
|
|
def encodes(self, x:Image.Image|TensorBBox|TensorPoint):
|
|
return x.crop_pad(self.size, self.tl, orig_sz=self.orig_sz)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #f0b73973
|
|
class OldRandomCrop(CropPad):
|
|
"Randomly crop an image to `size`"
|
|
def before_call(self, b, split_idx):
|
|
super().before_call(b, split_idx)
|
|
w,h = self.orig_sz
|
|
if not split_idx: self.tl = (random.randint(0,w-self.cp_size[0]), random.randint(0,h-self.cp_size[1]))
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #00527c2c
|
|
mk_class('ResizeMethod', **{o:o.lower() for o in ['Squish', 'Crop', 'Pad']},
|
|
doc="All possible resize method as attributes to get tab-completion and typo-proofing")
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #d93b6880
|
|
_all_ = ['ResizeMethod']
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #c1d598e8
|
|
@delegates()
|
|
class Resize(RandTransform):
|
|
split_idx,mode,mode_mask,order = None,BILINEAR,NEAREST,1
|
|
"Resize image to `size` using `method`"
|
|
def __init__(self,
|
|
size:int|tuple, # Size to resize to, duplicated if one value is specified
|
|
method:ResizeMethod=ResizeMethod.Crop, # A `ResizeMethod`
|
|
pad_mode:PadMode=PadMode.Reflection, # A `PadMode`
|
|
resamples=(BILINEAR, NEAREST), # Pillow `Image` resamples mode, resamples[1] for mask
|
|
**kwargs
|
|
):
|
|
size = _process_sz(size)
|
|
store_attr()
|
|
super().__init__(**kwargs)
|
|
self.mode,self.mode_mask = resamples
|
|
|
|
def before_call(self,
|
|
b,
|
|
split_idx:int # Index of the train/valid dataset
|
|
):
|
|
if self.method==ResizeMethod.Squish: return
|
|
self.pcts = (0.5,0.5) if split_idx else (random.random(),random.random())
|
|
|
|
def encodes(self, x:Image.Image|TensorBBox|TensorPoint):
|
|
orig_sz = _get_sz(x)
|
|
if self.method==ResizeMethod.Squish:
|
|
return x.crop_pad(orig_sz, fastuple(0,0), orig_sz=orig_sz, pad_mode=self.pad_mode,
|
|
resize_mode=self.mode_mask if isinstance(x,PILMask) else self.mode, resize_to=self.size)
|
|
|
|
w,h = orig_sz
|
|
op = (operator.lt,operator.gt)[self.method==ResizeMethod.Pad]
|
|
m = w/self.size[0] if op(w/self.size[0],h/self.size[1]) else h/self.size[1]
|
|
cp_sz = (int(m*self.size[0]),int(m*self.size[1]))
|
|
tl = fastuple(int(self.pcts[0]*(w-cp_sz[0])), int(self.pcts[1]*(h-cp_sz[1])))
|
|
return x.crop_pad(cp_sz, tl, orig_sz=orig_sz, pad_mode=self.pad_mode,
|
|
resize_mode=self.mode_mask if isinstance(x,PILMask) else self.mode, resize_to=self.size)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #c85221db
|
|
@delegates()
|
|
class RandomResizedCrop(RandTransform):
|
|
"Picks a random scaled crop of an image and resize it to `size`"
|
|
split_idx,order = None,1
|
|
def __init__(self,
|
|
size:int|tuple, # Final size, duplicated if one value is specified,,
|
|
min_scale:float=0.08, # Minimum scale of the crop, in relation to image area
|
|
ratio=(3/4, 4/3), # Range of width over height of the output
|
|
resamples=(BILINEAR, NEAREST), # Pillow `Image` resample mode, resamples[1] for mask
|
|
val_xtra:float=0.14, # The ratio of size at the edge cropped out in the validation set
|
|
max_scale:float=1., # Maximum scale of the crop, in relation to image area
|
|
**kwargs
|
|
):
|
|
size = _process_sz(size)
|
|
store_attr()
|
|
super().__init__(**kwargs)
|
|
self.mode,self.mode_mask = resamples
|
|
|
|
def before_call(self,
|
|
b,
|
|
split_idx # Index of the train/valid dataset
|
|
):
|
|
w,h = self.orig_sz = _get_sz(b)
|
|
if split_idx:
|
|
xtra = math.ceil(max(*self.size[:2])*self.val_xtra/8)*8
|
|
self.final_size = (self.size[0]+xtra, self.size[1]+xtra)
|
|
self.tl,self.cp_size = (0,0),self.orig_sz
|
|
return
|
|
self.final_size = self.size
|
|
for attempt in range(10):
|
|
area = random.uniform(self.min_scale, self.max_scale) * w * h
|
|
ratio = math.exp(random.uniform(math.log(self.ratio[0]), math.log(self.ratio[1])))
|
|
nw = int(round(math.sqrt(area * ratio)))
|
|
nh = int(round(math.sqrt(area / ratio)))
|
|
if nw <= w and nh <= h:
|
|
self.cp_size = (nw,nh)
|
|
self.tl = random.randint(0,w-nw), random.randint(0,h - nh)
|
|
return
|
|
if w/h < self.ratio[0]: self.cp_size = (w, int(w/self.ratio[0]))
|
|
elif w/h > self.ratio[1]: self.cp_size = (int(h*self.ratio[1]), h)
|
|
else: self.cp_size = (w, h)
|
|
self.tl = ((w-self.cp_size[0])//2, (h-self.cp_size[1])//2)
|
|
|
|
def encodes(self, x:Image.Image|TensorBBox|TensorPoint):
|
|
res = x.crop_pad(self.cp_size, self.tl, orig_sz=self.orig_sz,
|
|
resize_mode=self.mode_mask if isinstance(x,PILMask) else self.mode, resize_to=self.final_size)
|
|
if self.final_size != self.size: res = res.crop_pad(self.size) #Validation set: one final center crop
|
|
return res
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #3c1f51de
|
|
class RatioResize(DisplayedTransform):
|
|
'Resizes the biggest dimension of an image to `max_sz` maintaining the aspect ratio'
|
|
order = 1
|
|
def __init__(self,
|
|
max_sz: int, # Biggest dimension of the resized image
|
|
resamples=(BILINEAR, NEAREST), # Pillow `Image` resample mode, resamples[1] for mask
|
|
**kwargs
|
|
):
|
|
store_attr()
|
|
super().__init__(**kwargs)
|
|
|
|
def encodes(self, x:Image.Image|TensorBBox|TensorPoint):
|
|
w,h = _get_sz(x)
|
|
if w >= h: nw,nh = self.max_sz,h*self.max_sz/w
|
|
else: nw,nh = w*self.max_sz/h,self.max_sz
|
|
return Resize(size=(int(nh),int(nw)), resamples=self.resamples)(x)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #b1c9172b
|
|
def _init_mat(x):
|
|
mat = torch.eye(3, device=x.device).float()
|
|
return mat.unsqueeze(0).expand(x.size(0), 3, 3).contiguous()
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #98f63c80
|
|
def _grid_sample(x, coords, mode='bilinear', padding_mode='reflection', align_corners=None):
|
|
"Resample pixels in `coords` from `x` by `mode`, with `padding_mode` in ('reflection','border','zeros')."
|
|
#coords = coords.permute(0, 3, 1, 2).contiguous().permute(0, 2, 3, 1) # optimize layout for grid_sample
|
|
if mode=='bilinear': # hack to get smoother downwards resampling
|
|
mn,mx = coords.min(),coords.max()
|
|
# max amount we're affine zooming by (>1 means zooming in)
|
|
z = 1/(mx-mn).item()*2
|
|
# amount we're resizing by, with 100% extra margin
|
|
d = min(x.shape[-2]/coords.shape[-2], x.shape[-1]/coords.shape[-1])/2
|
|
# If we're resizing up by >200%, and we're zooming less than that, interpolate first
|
|
if d>1 and d>z:
|
|
x = F.interpolate(x, scale_factor=1/d, mode='area', recompute_scale_factor=True)
|
|
return F.grid_sample(x, coords, mode=mode, padding_mode=padding_mode, align_corners=align_corners)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #fca3bdd6
|
|
def affine_grid(
|
|
theta:Tensor, # Batch of affine transformation matrices
|
|
size:tuple, # Output size
|
|
align_corners:bool=None # PyTorch `F.grid_sample` align_corners
|
|
):
|
|
" Generates `TensorFlowField` from a transformation affine matrices `theta`"
|
|
return TensorFlowField(F.affine_grid(theta, size, align_corners=align_corners))
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #10b0ca25
|
|
@patch
|
|
def affine_coord(x: TensorImage,
|
|
mat:Tensor=None, # Batch of affine transformation matrices
|
|
coord_tfm:Callable=None, # Partial function of composable coordinate transforms
|
|
sz:int|tuple=None, # Output size, duplicated if one value is specified
|
|
mode:str='bilinear', # PyTorch `F.grid_sample` interpolation applied to `TensorImage`
|
|
pad_mode=PadMode.Reflection, # Padding applied to `TensorImage`
|
|
align_corners=True # PyTorch `F.grid_sample` align_corners
|
|
):
|
|
"Apply affine and coordinate transforms to `TensorImage`"
|
|
if mat is None and coord_tfm is None and sz is None: return x
|
|
size = tuple(x.shape[-2:]) if sz is None else (sz,sz) if isinstance(sz,int) else tuple(sz)
|
|
if mat is None: mat = _init_mat(x)[:,:2]
|
|
coords = affine_grid(mat, x.shape[:2] + size, align_corners=align_corners)
|
|
if coord_tfm is not None: coords = coord_tfm(coords)
|
|
return TensorImage(_grid_sample(x, coords, mode=mode, padding_mode=pad_mode, align_corners=align_corners))
|
|
|
|
@patch
|
|
def affine_coord(x: TensorMask,
|
|
mat:Tensor=None, # Batch of affine transformation matrices
|
|
coord_tfm:Callable=None, # Partial function of composable coordinate transforms
|
|
sz:int|tuple=None, # Output size, duplicated if one value is specified
|
|
mode='nearest', # PyTorch `F.grid_sample` interpolation applied to `TensorMask`
|
|
pad_mode=PadMode.Reflection, # Padding applied to `TensorMask`
|
|
align_corners=True # PyTorch `F.grid_sample` align_corners
|
|
):
|
|
"Apply affine and coordinate transforms to `TensorMask`"
|
|
add_dim = (x.ndim==3)
|
|
if add_dim: x = x[:,None]
|
|
res = TensorImage.affine_coord(x.float(), mat, coord_tfm, sz, mode, pad_mode, align_corners).long()
|
|
if add_dim: res = res[:,0]
|
|
return TensorMask(res)
|
|
|
|
@patch
|
|
def affine_coord(x: TensorPoint,
|
|
mat:Tensor=None, # Batch of affine transformation matrices
|
|
coord_tfm=None, # Partial function of composable coordinate transforms
|
|
sz=None, # Output size, duplicated if one value is specified
|
|
mode='nearest', # PyTorch `F.grid_sample` interpolation applied to `TensorPoint`
|
|
pad_mode=PadMode.Zeros, # Padding applied to `TensorPoint`
|
|
align_corners=True # PyTorch `F.grid_sample` align_corners
|
|
):
|
|
"Apply affine and coordinate transforms to `TensorPoint`"
|
|
#assert pad_mode==PadMode.Zeros, "Only zero padding is supported for `TensorPoint` and `TensorBBox`"
|
|
if sz is None: sz = getattr(x, "img_size", None)
|
|
if coord_tfm is not None: x = coord_tfm(x, invert=True)
|
|
if mat is not None:
|
|
mat = TensorPoint(mat)
|
|
x = (x - mat[:,:,2].unsqueeze(1)) @ torch.inverse(mat[:,:,:2].transpose(1,2))
|
|
return TensorPoint(x, sz=sz)
|
|
|
|
@patch
|
|
def affine_coord(x: TensorBBox,
|
|
mat=None, # Batch of affine transformation matrices
|
|
coord_tfm=None, # Partial function of composable coordinate transforms
|
|
sz=None, # Output size, duplicated if one value is specified
|
|
mode='nearest', # PyTorch `F.grid_sample` interpolation applied to `TensorBBox`
|
|
pad_mode=PadMode.Zeros, # Padding applied to `TensorBBox`
|
|
align_corners=True # PyTorch `F.grid_sample` align_corners
|
|
):
|
|
"Apply affine and coordinate transforms to `TensorBBox`"
|
|
if mat is None and coord_tfm is None: return x
|
|
if sz is None: sz = getattr(x, "img_size", None)
|
|
bs,n = x.shape[:2]
|
|
pnts = stack([x[...,:2], stack([x[...,0],x[...,3]],dim=2),
|
|
stack([x[...,2],x[...,1]],dim=2), x[...,2:]], dim=2)
|
|
pnts = TensorPoint(pnts.view(bs, 4*n, 2), img_size=sz).affine_coord(mat, coord_tfm, sz, mode, pad_mode)
|
|
pnts = pnts.view(bs, n, 4, 2)
|
|
tl,dr = pnts.min(dim=2)[0],pnts.max(dim=2)[0]
|
|
return TensorBBox(torch.cat([tl, dr], dim=2), img_size=sz)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #a1aeee8d
|
|
def _prepare_mat(x, mat):
|
|
h,w = getattr(x, 'img_size', x.shape[-2:])
|
|
mat[:,0,1] *= h/w
|
|
mat[:,1,0] *= w/h
|
|
return mat[:,:2]
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #64701200
|
|
class AffineCoordTfm(RandTransform):
|
|
"Combine and apply affine and coord transforms"
|
|
order,split_idx = 30,None
|
|
def __init__(self,
|
|
aff_fs:Callable|MutableSequence=None, # Affine transformations function for a batch
|
|
coord_fs:Callable|MutableSequence=None, # Coordinate transformations function for a batch
|
|
size:int|tuple=None, # Output size, duplicated if one value is specified
|
|
mode='bilinear', # PyTorch `F.grid_sample` interpolation
|
|
pad_mode=PadMode.Reflection, # A `PadMode`
|
|
mode_mask='nearest', # Resample mode for mask
|
|
align_corners=None, # PyTorch `F.grid_sample` align_corners
|
|
**kwargs
|
|
):
|
|
store_attr(but=['aff_fs','coord_fs'])
|
|
super().__init__(**kwargs)
|
|
self.aff_fs,self.coord_fs = L(aff_fs),L(coord_fs)
|
|
self.cp_size = None if size is None else (size,size) if isinstance(size, int) else tuple(size)
|
|
|
|
def before_call(self,
|
|
b,
|
|
split_idx, # Index of the train/valid dataset
|
|
):
|
|
while isinstance(b, tuple): b = b[0]
|
|
self.split_idx = split_idx
|
|
self.do,self.mat = True,self._get_affine_mat(b)
|
|
for t in self.coord_fs: t.before_call(b)
|
|
|
|
def compose(self, tfm):
|
|
"Compose `self` with another `AffineCoordTfm` to only do the interpolation step once"
|
|
# TODO: keep `name` up to date with the combination
|
|
# TODO: have option to only show a subset of the attrs, e.g. for `Flip`
|
|
self.aff_fs += tfm.aff_fs
|
|
self.coord_fs += tfm.coord_fs
|
|
|
|
def _get_affine_mat(self, x):
|
|
aff_m = _init_mat(x)
|
|
if self.split_idx: return _prepare_mat(x, aff_m)
|
|
ms = [f(x) for f in self.aff_fs]
|
|
ms = [m for m in ms if m is not None]
|
|
for m in ms: aff_m = aff_m @ m
|
|
return _prepare_mat(x, aff_m)
|
|
|
|
def _encode(self, x, mode, reverse=False):
|
|
coord_func = None if len(self.coord_fs)==0 or self.split_idx else partial(compose_tfms, tfms=self.coord_fs, reverse=reverse)
|
|
return x.affine_coord(self.mat, coord_func, sz=self.size, mode=mode, pad_mode=self.pad_mode, align_corners=self.align_corners)
|
|
|
|
def encodes(self, x:TensorImage): return self._encode(x, self.mode)
|
|
def encodes(self, x:TensorMask): return self._encode(x, self.mode_mask)
|
|
def encodes(self, x:TensorPoint|TensorBBox): return self._encode(x, self.mode, reverse=True)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #259917d1
|
|
class RandomResizedCropGPU(RandTransform):
|
|
"Picks a random scaled crop of an image and resize it to `size`"
|
|
split_idx,order = None,30
|
|
def __init__(self,
|
|
size, # Final size, duplicated if one value is specified
|
|
min_scale=0.08, # Minimum scale of the crop, in relation to image area
|
|
ratio=(3/4, 4/3), # Range of width over height of the output
|
|
mode='bilinear', # PyTorch `F.grid_sample` interpolation
|
|
valid_scale=1., # Scale of the crop for the validation set, in relation to image area
|
|
max_scale=1., # Maximum scale of the crop, in relation to image area
|
|
mode_mask='nearest', # Interpolation mode for `TensorMask`
|
|
**kwargs
|
|
):
|
|
if isinstance(size, int): size = (size,size)
|
|
store_attr()
|
|
super().__init__(**kwargs)
|
|
|
|
def before_call(self, b, split_idx):
|
|
self.do = True
|
|
h,w = fastuple((b[0] if isinstance(b, tuple) else b).shape[-2:])
|
|
for attempt in range(10):
|
|
if split_idx: break
|
|
area = random.uniform(self.min_scale,self.max_scale) * w * h
|
|
ratio = math.exp(random.uniform(math.log(self.ratio[0]), math.log(self.ratio[1])))
|
|
nw = int(round(math.sqrt(area * ratio)))
|
|
nh = int(round(math.sqrt(area / ratio)))
|
|
if nw <= w and nh <= h:
|
|
self.cp_size = (nh,nw)
|
|
self.tl = random.randint(0,h - nh),random.randint(0,w-nw)
|
|
return
|
|
if w/h < self.ratio[0]: self.cp_size = (int(w/self.ratio[0]), w)
|
|
elif w/h > self.ratio[1]: self.cp_size = (h, int(h*self.ratio[1]))
|
|
else: self.cp_size = (h, w)
|
|
if split_idx: self.cp_size = (int(self.cp_size[0]*self.valid_scale), int(self.cp_size[1]*self.valid_scale))
|
|
self.tl = ((h-self.cp_size[0])//2,(w-self.cp_size[1])//2)
|
|
|
|
def _encode(self, x, mode):
|
|
x = x[...,self.tl[0]:self.tl[0]+self.cp_size[0], self.tl[1]:self.tl[1]+self.cp_size[1]]
|
|
return x.affine_coord(sz=self.size, mode=mode)
|
|
|
|
def encodes(self, x:TensorImage|TensorPoint|TensorBBox): return self._encode(x, self.mode)
|
|
def encodes(self, x:TensorMask): return self._encode(x, self.mode_mask)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #879b316d
|
|
def mask_tensor(
|
|
x:Tensor, # Input `Tensor`
|
|
p=0.5, # Probability of not applying mask
|
|
neutral=0., # Mask value
|
|
batch=False # Apply identical mask to entire batch
|
|
):
|
|
"Mask elements of `x` with `neutral` with probability `1-p`"
|
|
if p==1.: return x
|
|
if batch: return x if random.random() < p else x.new_zeros(*x.size()) + neutral
|
|
if neutral != 0: x.add_(-neutral)
|
|
# Extra casting to float and long to prevent crashes on mps accelerator (issue #3911)
|
|
mask = x.new_empty(*x.size()).float().bernoulli_(p).long()
|
|
x.mul_(mask)
|
|
return x.add_(neutral) if neutral != 0 else x
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #ad5b7f65
|
|
def _draw_mask(x, def_draw, draw=None, p=0.5, neutral=0., batch=False):
|
|
"Creates mask_tensor based on `x` with `neutral` with probability `1-p`. "
|
|
if draw is None: draw=def_draw
|
|
if callable(draw): res=draw(x)
|
|
elif is_listy(draw):
|
|
assert len(draw)>=x.size(0)
|
|
res = tensor(draw[:x.size(0)], dtype=x.dtype, device=x.device)
|
|
else: res = x.new_zeros(x.size(0)) + draw
|
|
return TensorBase(mask_tensor(res, p=p, neutral=neutral, batch=batch))
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #4cc16090
|
|
def affine_mat(*ms):
|
|
"Restructure length-6 vector `ms` into an affine matrix with 0,0,1 in the last line"
|
|
return stack([stack([ms[0], ms[1], ms[2]], dim=1),
|
|
stack([ms[3], ms[4], ms[5]], dim=1),
|
|
stack([t0(ms[0]), t0(ms[0]), t1(ms[0])], dim=1)], dim=1)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #0b648343
|
|
def flip_mat(
|
|
x:Tensor, # The input Tensor
|
|
p=0.5, # Probability of appying transformation
|
|
draw:int|MutableSequence|Callable=None, # Custom flips instead of random
|
|
batch:bool=False # Apply identical flip to entire batch
|
|
):
|
|
"Return a random flip matrix"
|
|
def _def_draw(x): return x.new_ones(x.size(0))
|
|
mask = x.new_ones(x.size(0)) - 2*_draw_mask(x, _def_draw, draw=draw, p=p, batch=batch)
|
|
return affine_mat(mask, t0(mask), t0(mask),
|
|
t0(mask), t1(mask), t0(mask))
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #ca49f33a
|
|
def _get_default(x, mode=None, pad_mode=None):
|
|
if mode is None: mode='bilinear' if isinstance(x, TensorMask) else 'bilinear'
|
|
if pad_mode is None: pad_mode=PadMode.Zeros if isinstance(x, (TensorPoint, TensorBBox)) else PadMode.Reflection
|
|
x0 = x[0] if isinstance(x, tuple) else x
|
|
return x0,mode,pad_mode
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #583c0e30
|
|
@patch
|
|
def flip_batch(x: TensorImage|TensorMask|TensorPoint|TensorBBox,
|
|
p=0.5, # Probability of applying flip
|
|
draw:int|MutableSequence|Callable=None, # Custom flips instead of random
|
|
size:int|tuple=None, # Output size, duplicated if one value is specified
|
|
mode=None, # PyTorch `F.grid_sample` interpolation applied to `x`
|
|
pad_mode=None, # Padding applied to `x`
|
|
align_corners=True, # PyTorch `F.grid_sample` align_corners
|
|
batch=False # Apply identical flip to entire batch
|
|
):
|
|
x0,mode,pad_mode = _get_default(x, mode, pad_mode)
|
|
mat=flip_mat(x0, p=p, draw=draw, batch=batch)
|
|
return x.affine_coord(mat=mat[:,:2], sz=size, mode=mode, pad_mode=pad_mode, align_corners=align_corners)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #0979aa4f
|
|
class Flip(AffineCoordTfm):
|
|
"Randomly flip a batch of images with a probability `p`"
|
|
def __init__(self,
|
|
p=0.5, # Probability of applying flip
|
|
draw:int|MutableSequence|Callable=None, # Custom flips instead of random
|
|
size:int|tuple=None, # Output size, duplicated if one value is specified
|
|
mode:str='bilinear', # PyTorch `F.grid_sample` interpolation
|
|
pad_mode=PadMode.Reflection, # A `PadMode`
|
|
align_corners=True, # PyTorch `F.grid_sample` align_corners
|
|
batch=False # Apply identical flip to entire batch
|
|
):
|
|
aff_fs = partial(flip_mat, p=p, draw=draw, batch=batch)
|
|
super().__init__(aff_fs, size=size, mode=mode, pad_mode=pad_mode, align_corners=align_corners, p=p)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #74aa2c86
|
|
class DeterministicDraw():
|
|
def __init__(self, vals): self.vals,self.count = vals,-1
|
|
|
|
def __call__(self, x):
|
|
self.count += 1
|
|
return x.new_zeros(x.size(0)) + self.vals[self.count%len(self.vals)]
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #22032fee
|
|
class DeterministicFlip(Flip):
|
|
"Flip the batch every other call"
|
|
def __init__(self,
|
|
size:int|tuple=None, # Output size, duplicated if one value is specified
|
|
mode:str='bilinear', # PyTorch `F.grid_sample` interpolation
|
|
pad_mode=PadMode.Reflection, # A `PadMode`
|
|
align_corners=True, # PyTorch `F.grid_sample` align_corners
|
|
**kwargs
|
|
):
|
|
super().__init__(p=1., draw=DeterministicDraw([0,1]), mode=mode, pad_mode=pad_mode, align_corners=align_corners, **kwargs)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #80222ddc
|
|
def dihedral_mat(
|
|
x:Tensor, # Input `Tensor`
|
|
p:float=0.5, # Probability of staying unchanged
|
|
draw:int|MutableSequence|Callable=None, # Custom dihedrals instead of random
|
|
batch:bool=False # Apply identical dihedral to entire batch
|
|
):
|
|
"Return a random dihedral matrix"
|
|
def _def_draw(x): return torch.randint(0,8, (x.size(0),), device=x.device)
|
|
def _def_draw_b(x): return random.randint(0,7) + x.new_zeros((x.size(0),)).long()
|
|
idx = _draw_mask(x, _def_draw_b if batch else _def_draw, draw=draw, p=p, batch=batch).long()
|
|
xs = tensor([1,-1,1,-1,-1,1,1,-1], device=x.device).gather(0, idx)
|
|
ys = tensor([1,1,-1,1,-1,-1,1,-1], device=x.device).gather(0, idx)
|
|
m0 = tensor([1,1,1,0,1,0,0,0], device=x.device).gather(0, idx)
|
|
m1 = tensor([0,0,0,1,0,1,1,1], device=x.device).gather(0, idx)
|
|
return affine_mat(xs*m0, xs*m1, t0(xs),
|
|
ys*m1, ys*m0, t0(xs)).float()
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #3402ff49
|
|
@patch
|
|
def dihedral_batch(x: TensorImage|TensorMask|TensorPoint|TensorBBox,
|
|
p=0.5, # Probability of applying dihedral
|
|
draw:int|MutableSequence|Callable=None, # Custom dihedrals instead of random
|
|
size:int|tuple=None, # Output size, duplicated if one value is specified
|
|
mode:str='bilinear', # PyTorch `F.grid_sample` interpolation applied to `x`
|
|
pad_mode=None, # Padding applied to `x`
|
|
batch=False, # Apply identical dihedral to entire batch
|
|
align_corners=True # PyTorch `F.grid_sample` align_corners
|
|
):
|
|
x0,mode,pad_mode = _get_default(x, mode, pad_mode)
|
|
mat = _prepare_mat(x, dihedral_mat(x0, p=p, draw=draw, batch=batch))
|
|
return x.affine_coord(mat=mat, sz=size, mode=mode, pad_mode=pad_mode, align_corners=align_corners)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #a2c35d51
|
|
class Dihedral(AffineCoordTfm):
|
|
"Apply a random dihedral transformation to a batch of images with a probability `p`"
|
|
def __init__(self,
|
|
p=0.5, # Probability of applying dihedral
|
|
draw:int|MutableSequence|Callable=None, # Custom dihedrals instead of random
|
|
size:int|tuple=None, # Output size, duplicated if one value is specified
|
|
mode:str='bilinear', # PyTorch `F.grid_sample` interpolation
|
|
pad_mode=PadMode.Reflection, # A `PadMode`
|
|
batch=False, # Apply identical dihedral to entire batch
|
|
align_corners=True # PyTorch `F.grid_sample` align_corners
|
|
):
|
|
f = partial(dihedral_mat, p=p, draw=draw, batch=batch)
|
|
super().__init__(aff_fs=f, size=size, mode=mode, pad_mode=pad_mode, align_corners=align_corners)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #41f12939
|
|
class DeterministicDihedral(Dihedral):
|
|
def __init__(self,
|
|
size:int|tuple=None, # Output size, duplicated if one value is specified
|
|
mode:str='bilinear', # PyTorch `F.grid_sample` interpolation
|
|
pad_mode=PadMode.Reflection, # A `PadMode`
|
|
align_corners=None # PyTorch `F.grid_sample` align_corners
|
|
):
|
|
"Flip the batch every other call"
|
|
super().__init__(p=1., draw=DeterministicDraw(list(range(8))), pad_mode=pad_mode, align_corners=align_corners)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #abb5081d
|
|
def rotate_mat(
|
|
x:Tensor, # Input `Tensor`
|
|
max_deg:int=10, # Maximum degree of rotation
|
|
p:float=0.5, # Probability of applying rotate
|
|
draw:int|MutableSequence|Callable=None, # Custom rotates instead of random
|
|
batch:bool=False # Apply identical rotate to entire batch
|
|
):
|
|
"Return a random rotation matrix with `max_deg` and `p`"
|
|
def _def_draw(x): return x.new_empty(x.size(0)).uniform_(-max_deg, max_deg)
|
|
def _def_draw_b(x): return x.new_zeros(x.size(0)) + random.uniform(-max_deg, max_deg)
|
|
thetas = _draw_mask(x, _def_draw_b if batch else _def_draw, draw=draw, p=p, batch=batch) * math.pi/180
|
|
return affine_mat(thetas.cos(), thetas.sin(), t0(thetas),
|
|
-thetas.sin(), thetas.cos(), t0(thetas))
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #813ae11c
|
|
@patch
|
|
@delegates(rotate_mat)
|
|
def rotate(x: TensorImage|TensorMask|TensorPoint|TensorBBox,
|
|
size:int|tuple=None, # Output size, duplicated if one value is specified
|
|
mode:str=None, # PyTorch `F.grid_sample` interpolation applied to `x`
|
|
pad_mode=None, # Padding applied to `x`
|
|
align_corners:bool=True, # PyTorch `F.grid_sample` align_corners
|
|
**kwargs
|
|
):
|
|
x0,mode,pad_mode = _get_default(x, mode, pad_mode)
|
|
mat = _prepare_mat(x, rotate_mat(x0, **kwargs))
|
|
return x.affine_coord(mat=mat, sz=size, mode=mode, pad_mode=pad_mode, align_corners=align_corners)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #6de3f169
|
|
class Rotate(AffineCoordTfm):
|
|
"Apply a random rotation of at most `max_deg` with probability `p` to a batch of images"
|
|
def __init__(self,
|
|
max_deg:int=10, # Maximum degree of rotation
|
|
p:float=0.5, # Probability of applying rotate
|
|
draw:int|MutableSequence|Callable=None, # Custom rotates instead of random
|
|
size:int|tuple=None, # Output size, duplicated if one value is specified
|
|
mode:str='bilinear', # PyTorch `F.grid_sample` interpolation
|
|
pad_mode=PadMode.Reflection, # A `PadMode`
|
|
align_corners:bool=True, # PyTorch `F.grid_sample` align_corners
|
|
batch:bool=False # Apply identical rotate to entire batch
|
|
):
|
|
aff_fs = partial(rotate_mat, max_deg=max_deg, p=p, draw=draw, batch=batch)
|
|
super().__init__(aff_fs=aff_fs, size=size, mode=mode, pad_mode=pad_mode, align_corners=align_corners)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #911eea62
|
|
def zoom_mat(
|
|
x:Tensor, # Input `Tensor`
|
|
min_zoom:float=1., # Minimum zoom
|
|
max_zoom:float=1.1, # Maximum zoom
|
|
p:float=0.5, # Probability of applying zoom
|
|
draw:float|MutableSequence|Callable=None, # User defined scale of the zoom
|
|
draw_x:float|MutableSequence|Callable=None, # User defined center of the zoom in x
|
|
draw_y:float|MutableSequence|Callable=None, # User defined center of the zoom in y
|
|
batch:bool=False # Apply identical zoom to entire batch
|
|
):
|
|
"Return a random zoom matrix with `max_zoom` and `p`"
|
|
def _def_draw(x): return x.new_empty(x.size(0)).uniform_(min_zoom, max_zoom)
|
|
def _def_draw_b(x): return x.new_zeros(x.size(0)) + random.uniform(min_zoom, max_zoom)
|
|
def _def_draw_ctr(x): return x.new_empty(x.size(0)).uniform_(0,1)
|
|
def _def_draw_ctr_b(x): return x.new_zeros(x.size(0)) + random.uniform(0,1)
|
|
assert(min_zoom<=max_zoom)
|
|
s = 1/_draw_mask(x, _def_draw_b if batch else _def_draw, draw=draw, p=p, neutral=1., batch=batch)
|
|
def_draw_c = _def_draw_ctr_b if batch else _def_draw_ctr
|
|
col_pct = _draw_mask(x, def_draw_c, draw=draw_x, p=1., batch=batch)
|
|
row_pct = _draw_mask(x, def_draw_c, draw=draw_y, p=1., batch=batch)
|
|
col_c = (1-s) * (2*col_pct - 1)
|
|
row_c = (1-s) * (2*row_pct - 1)
|
|
return affine_mat(s, t0(s), col_c,
|
|
t0(s), s, row_c)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #76ecc04c
|
|
@patch
|
|
@delegates(zoom_mat)
|
|
def zoom(x: TensorImage|TensorMask|TensorPoint|TensorBBox,
|
|
size:int|tuple=None, # Output size, duplicated if one value is specified
|
|
mode:str='bilinear', # PyTorch `F.grid_sample` interpolation applied to `x`
|
|
pad_mode=PadMode.Reflection, # Padding applied to `x`
|
|
align_corners:bool=True, # PyTorch `F.grid_sample` align_corners
|
|
**kwargs
|
|
):
|
|
x0,mode,pad_mode = _get_default(x, mode, pad_mode)
|
|
return x.affine_coord(mat=zoom_mat(x0, **kwargs)[:,:2], sz=size, mode=mode, pad_mode=pad_mode, align_corners=align_corners)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #10be7308
|
|
class Zoom(AffineCoordTfm):
|
|
"Apply a random zoom of at most `max_zoom` with probability `p` to a batch of images"
|
|
def __init__(self,
|
|
min_zoom:float=1., # Minimum zoom
|
|
max_zoom:float=1.1, # Maximum zoom
|
|
p:float=0.5, # Probability of applying zoom
|
|
draw:float|MutableSequence|Callable=None, # User defined scale of the zoom
|
|
draw_x:float|MutableSequence|Callable=None, # User defined center of the zoom in x
|
|
draw_y:float|MutableSequence|Callable=None, # User defined center of the zoom in y
|
|
size:int|tuple=None, # Output size, duplicated if one value is specified
|
|
mode='bilinear', # PyTorch `F.grid_sample` interpolation
|
|
pad_mode=PadMode.Reflection, # A `PadMode`
|
|
batch=False, # Apply identical zoom to entire batch
|
|
align_corners=True # PyTorch `F.grid_sample` align_corners
|
|
):
|
|
aff_fs = partial(zoom_mat, min_zoom=min_zoom, max_zoom=max_zoom, p=p, draw=draw, draw_x=draw_x, draw_y=draw_y, batch=batch)
|
|
super().__init__(aff_fs, size=size, mode=mode, pad_mode=pad_mode, align_corners=align_corners)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #f59ebe09
|
|
def solve(A,B):
|
|
return torch.linalg.solve(A,B)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #dad4e52a
|
|
def find_coeffs(
|
|
p1:Tensor, # Original points
|
|
p2:Tensor, # Target points
|
|
):
|
|
"Find coefficients for warp tfm from `p1` to `p2`"
|
|
m = []
|
|
p = p1[:,0,0]
|
|
#The equations we'll need to solve.
|
|
for i in range(p1.shape[1]):
|
|
m.append(stack([p2[:,i,0], p2[:,i,1], t1(p), t0(p), t0(p), t0(p), -p1[:,i,0]*p2[:,i,0], -p1[:,i,0]*p2[:,i,1]]))
|
|
m.append(stack([t0(p), t0(p), t0(p), p2[:,i,0], p2[:,i,1], t1(p), -p1[:,i,1]*p2[:,i,0], -p1[:,i,1]*p2[:,i,1]]))
|
|
#The 8 scalars we seek are solution of AX = B
|
|
A = stack(m).permute(2, 0, 1)
|
|
B = p1.view(p1.shape[0], 8, 1)
|
|
return solve(A,B)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #17d093aa
|
|
def apply_perspective(
|
|
coords:Tensor, # Original coordinates
|
|
coeffs:Tensor # Warping transformation matrice
|
|
):
|
|
"Apply perspective tranform on `coords` with `coeffs`"
|
|
sz = coords.shape
|
|
coords = coords.view(sz[0], -1, 2)
|
|
coeffs = torch.cat([coeffs, t1(coeffs[:,:1])], dim=1).view(coeffs.shape[0], 3,3)
|
|
coords1 = coords @ coeffs[...,:2].transpose(1,2) + coeffs[...,2].unsqueeze(1)
|
|
if (coords1[...,2]==0.).any(): return coords[...,:2].view(*sz)
|
|
coords = coords1/coords1[...,2].unsqueeze(-1)
|
|
return coords[...,:2].view(*sz)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #74394fdb
|
|
class _WarpCoord():
|
|
def __init__(self, magnitude=0.2, p=0.5, draw_x=None, draw_y=None, batch=False):
|
|
store_attr()
|
|
self.coeffs = None
|
|
|
|
def _def_draw(self, x):
|
|
if not self.batch: return x.new_empty(x.size(0)).uniform_(-self.magnitude, self.magnitude)
|
|
return x.new_zeros(x.size(0)) + random.uniform(-self.magnitude, self.magnitude)
|
|
|
|
def before_call(self, x):
|
|
x_t = _draw_mask(x, self._def_draw, self.draw_x, p=self.p, batch=self.batch)
|
|
y_t = _draw_mask(x, self._def_draw, self.draw_y, p=self.p, batch=self.batch)
|
|
orig_pts = torch.tensor([[-1,-1], [-1,1], [1,-1], [1,1]], dtype=x.dtype, device=x.device)
|
|
self.orig_pts = orig_pts.unsqueeze(0).expand(x.size(0),4,2)
|
|
targ_pts = stack([stack([-1-y_t, -1-x_t]), stack([-1+y_t, 1+x_t]),
|
|
stack([ 1+y_t, -1+x_t]), stack([ 1-y_t, 1-x_t])])
|
|
self.targ_pts = targ_pts.permute(2,0,1)
|
|
|
|
def __call__(self, x, invert=False):
|
|
coeffs = find_coeffs(self.targ_pts, self.orig_pts) if invert else find_coeffs(self.orig_pts, self.targ_pts)
|
|
return apply_perspective(x, coeffs)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #6014cd7f
|
|
@patch
|
|
@delegates(_WarpCoord.__init__)
|
|
def warp(x:TensorImage|TensorMask|TensorPoint|TensorBBox,
|
|
size:int|tuple=None, # Output size, duplicated if one value is specified
|
|
mode:str='bilinear', # PyTorch `F.grid_sample` interpolation applied to `x`
|
|
pad_mode=PadMode.Reflection, # Padding applied to `x`
|
|
align_corners:bool=True, # PyTorch `F.grid_sample` align_corners
|
|
**kwargs
|
|
):
|
|
x0,mode,pad_mode = _get_default(x, mode, pad_mode)
|
|
coord_tfm = _WarpCoord(**kwargs)
|
|
coord_tfm.before_call(x0)
|
|
return x.affine_coord(coord_tfm=coord_tfm, sz=size, mode=mode, pad_mode=pad_mode, align_corners=align_corners)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #a11ca7c1
|
|
class Warp(AffineCoordTfm):
|
|
"Apply perspective warping with `magnitude` and `p` on a batch of matrices"
|
|
def __init__(self,
|
|
magnitude:float=0.2, # The default warping magnitude
|
|
p:float=0.5, # Probability of applying warp
|
|
draw_x:float|MutableSequence|Callable=None, # User defined warping magnitude in x
|
|
draw_y:float|MutableSequence|Callable=None, # User defined warping magnitude in y
|
|
size:int|tuple=None, # Output size, duplicated if one value is specified
|
|
mode:str='bilinear', # PyTorch `F.grid_sample` interpolation
|
|
pad_mode=PadMode.Reflection, # A `PadMode`
|
|
batch:bool=False, # Apply identical warp to entire batch
|
|
align_corners:bool=True # PyTorch `F.grid_sample` align_corners
|
|
):
|
|
store_attr()
|
|
coord_fs = _WarpCoord(magnitude=magnitude, p=p, draw_x=draw_x, draw_y=draw_y, batch=batch)
|
|
super().__init__(coord_fs=coord_fs, size=size, mode=mode, pad_mode=pad_mode, align_corners=align_corners )
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #b7a0cc40
|
|
@patch
|
|
def lighting(x: TensorImage, func): return torch.sigmoid(func(logit(x)))
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #b2dfa28c
|
|
class SpaceTfm(RandTransform):
|
|
"Apply `fs` to the logits"
|
|
order = 40
|
|
def __init__(self,
|
|
fs:Callable|MutableSequence, # Transformation functions applying in a space
|
|
space_fn:Callable, # Function converting rgb to a space and back to rgb after appying `fs`
|
|
**kwargs
|
|
):
|
|
super().__init__(**kwargs)
|
|
self.space_fn=space_fn
|
|
self.fs=L(fs)
|
|
|
|
def before_call(self,
|
|
b,
|
|
split_idx:int, # Index of the train/valid dataset
|
|
):
|
|
self.do = True
|
|
while isinstance(b, tuple): b = b[0]
|
|
for t in self.fs: t.before_call(b)
|
|
|
|
def compose(self,
|
|
tfm:Callable # Transformation function to compose
|
|
):
|
|
"Compose `self` with another `LightingTransform`"
|
|
self.fs += tfm.fs
|
|
|
|
def encodes(self,x:TensorImage): return self.space_fn(x,partial(compose_tfms, tfms=self.fs))
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #8b65472d
|
|
class LightingTfm(SpaceTfm):
|
|
"Apply `fs` to the logits"
|
|
order = 40
|
|
def __init__(self,
|
|
fs:Callable|MutableSequence, # Transformation functions applying in logit space,
|
|
**kwargs
|
|
):
|
|
super().__init__(fs, TensorImage.lighting, **kwargs)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #2ca4c6dd
|
|
class _BrightnessLogit():
|
|
def __init__(self, max_lighting=0.2, p=0.75, draw=None, batch=False): store_attr()
|
|
|
|
def _def_draw(self, x):
|
|
if not self.batch: return x.new_empty(x.size(0)).uniform_(0.5*(1-self.max_lighting), 0.5*(1+self.max_lighting))
|
|
return x.new_zeros(x.size(0)) + random.uniform(0.5*(1-self.max_lighting), 0.5*(1+self.max_lighting))
|
|
|
|
def before_call(self, x):
|
|
self.change = _draw_mask(x, self._def_draw, draw=self.draw, p=self.p, neutral=0.5, batch=self.batch)
|
|
|
|
def __call__(self, x): return x.add_(logit(self.change[:,None,None,None]))
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #82bc73f8
|
|
@patch
|
|
@delegates(_BrightnessLogit.__init__)
|
|
def brightness(x: TensorImage, **kwargs):
|
|
func = _BrightnessLogit(**kwargs)
|
|
func.before_call(x)
|
|
return x.lighting(func)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #4b321b7a
|
|
class Brightness(LightingTfm):
|
|
def __init__(self,
|
|
max_lighting:float=0.2, # Maximum scale of changing brightness
|
|
p:float=0.75, # Probability of appying transformation
|
|
draw:float|MutableSequence|Callable=None, # User defined behavior of batch transformation
|
|
batch=False # Apply identical brightness to entire batch
|
|
):
|
|
"Apply change in brightness of `max_lighting` to batch of images with probability `p`."
|
|
store_attr()
|
|
super().__init__(_BrightnessLogit(max_lighting, p, draw, batch))
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #db1cb711
|
|
class _ContrastLogit():
|
|
def __init__(self, max_lighting=0.2, p=0.75, draw=None, batch=False): store_attr()
|
|
|
|
def _def_draw(self, x):
|
|
if not self.batch: res = x.new_empty(x.size(0)).uniform_(math.log(1-self.max_lighting), -math.log(1-self.max_lighting))
|
|
else: res = x.new_zeros(x.size(0)) + random.uniform(math.log(1-self.max_lighting), -math.log(1-self.max_lighting))
|
|
return torch.exp(res)
|
|
|
|
def before_call(self, x):
|
|
self.change = _draw_mask(x, self._def_draw, draw=self.draw, p=self.p, neutral=1., batch=self.batch)
|
|
|
|
def __call__(self, x): return x.mul_(self.change[:,None,None,None])
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #01929d72
|
|
@patch
|
|
@delegates(_ContrastLogit.__init__)
|
|
def contrast(x: TensorImage, **kwargs):
|
|
func = _ContrastLogit(**kwargs)
|
|
func.before_call(x)
|
|
return x.lighting(func)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #c08daacc
|
|
class Contrast(LightingTfm):
|
|
"Apply change in contrast of `max_lighting` to batch of images with probability `p`."
|
|
def __init__(self,
|
|
max_lighting=0.2, # Maximum scale of changing contrast
|
|
p=0.75, # Probability of appying transformation
|
|
draw:float|MutableSequence|Callable=None, # User defined behavior of batch transformation
|
|
batch=False
|
|
):
|
|
store_attr()
|
|
super().__init__(_ContrastLogit(max_lighting, p, draw, batch))
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #2cc9a757
|
|
def grayscale(x):
|
|
"Tensor to grayscale tensor. Uses the ITU-R 601-2 luma transform. "
|
|
return (x*torch.tensor([0.2989,0.5870,0.1140],device=x.device)[...,None,None]).sum(1)[:,None]
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #52e50a03
|
|
class _SaturationLogit():
|
|
def __init__(self, max_lighting=0.2, p=0.75, draw=None, batch=False): store_attr()
|
|
|
|
def _def_draw(self, x):
|
|
if not self.batch: res = x.new_empty(x.size(0)).uniform_(math.log(1-self.max_lighting), -math.log(1-self.max_lighting))
|
|
else: res = x.new_zeros(x.size(0)) + random.uniform(math.log(1-self.max_lighting), -math.log(1-self.max_lighting))
|
|
return torch.exp(res)
|
|
|
|
def before_call(self, x):
|
|
self.change = _draw_mask(x, self._def_draw, draw=self.draw, p=self.p, neutral=1., batch=self.batch)
|
|
|
|
def __call__(self, x):
|
|
#interpolate between grayscale and original in-place
|
|
gs = grayscale(x)
|
|
gs.mul_(1-self.change[:,None,None,None])
|
|
x.mul_(self.change[:,None,None,None])
|
|
return x.add_(gs)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #e5903c98
|
|
@patch
|
|
@delegates(_SaturationLogit.__init__)
|
|
def saturation(x: TensorImage, **kwargs):
|
|
func = _SaturationLogit(**kwargs)
|
|
func.before_call(x)
|
|
return x.lighting(func)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #b8d9f854
|
|
class Saturation(LightingTfm):
|
|
"Apply change in saturation of `max_lighting` to batch of images with probability `p`."
|
|
# Ref: https://pytorch.org/docs/stable/torchvision/transforms.html#torchvision.transforms.functional.adjust_saturation
|
|
def __init__(self,
|
|
max_lighting:float=0.2, # Maximum scale of changing brightness
|
|
p:float=0.75, # Probability of appying transformation
|
|
draw:float|MutableSequence|Callable=None, # User defined behavior of batch transformation
|
|
batch:bool=False # Apply identical saturation to entire batch
|
|
):
|
|
store_attr()
|
|
super().__init__(_SaturationLogit(max_lighting, p, draw, batch))
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #e9d30c93
|
|
def rgb2hsv(
|
|
img:Tensor # Batch of images `Tensor`in RGB
|
|
):
|
|
"Converts a RGB image to an HSV image. Note: Will not work on logit space images."
|
|
r, g, b = img.unbind(1)
|
|
# temp commented out due to https://github.com/pytorch/pytorch/issues/47069
|
|
# maxc = torch.max(img, dim=1).values
|
|
# minc = torch.min(img, dim=1).values
|
|
maxc = torch.max(img, dim=1)[0]
|
|
minc = torch.min(img, dim=1)[0]
|
|
eqc = maxc == minc
|
|
|
|
cr = maxc - minc
|
|
s = cr / torch.where(eqc, maxc.new_ones(()), maxc)
|
|
cr_divisor = torch.where(eqc, maxc.new_ones(()), cr)
|
|
rc = (maxc - r) / cr_divisor
|
|
gc = (maxc - g) / cr_divisor
|
|
bc = (maxc - b) / cr_divisor
|
|
|
|
hr = (maxc == r) * (bc - gc)
|
|
hg = ((maxc == g) & (maxc != r)) * (2.0 + rc - bc)
|
|
hb = ((maxc != g) & (maxc != r)) * (4.0 + gc - rc)
|
|
h = (hr + hg + hb)
|
|
h = torch.fmod((h / 6.0 + 1.0), 1.0)
|
|
|
|
return torch.stack((h, s, maxc),dim=1)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #42b94e0c
|
|
def hsv2rgb(
|
|
img:Tensor, # Batch of images `Tensor in HSV`
|
|
):
|
|
"Converts a HSV image to an RGB image."
|
|
h, s, v = img.unbind(1)
|
|
i = torch.floor(h * 6.0)
|
|
f = (h * 6.0) - i
|
|
i = i.to(dtype=torch.int32)
|
|
|
|
p = torch.clamp((v * (1.0 - s)), 0.0, 1.0)
|
|
q = torch.clamp((v * (1.0 - s * f)), 0.0, 1.0)
|
|
t = torch.clamp((v * (1.0 - s * (1.0 - f))), 0.0, 1.0)
|
|
i = i % 6
|
|
|
|
mask = i[:,None] == torch.arange(6,device=i.device)[:, None, None][None]
|
|
|
|
a1 = torch.stack((v, q, p, p, t, v),dim=1)
|
|
a2 = torch.stack((t, v, v, q, p, p),dim=1)
|
|
a3 = torch.stack((p, p, t, v, v, q),dim=1)
|
|
a4 = torch.stack((a1, a2, a3),dim=1)
|
|
|
|
return torch.einsum("nijk, nxijk -> nxjk", mask.to(dtype=img.dtype), a4)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #9e74e96b
|
|
@patch
|
|
def hsv(x: TensorImage, func): return TensorImage(hsv2rgb(func(rgb2hsv(x))))
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #8c609dc6
|
|
class HSVTfm(SpaceTfm):
|
|
"Apply `fs` to the images in HSV space"
|
|
def __init__(self, fs, **kwargs):
|
|
super().__init__(fs, TensorImage.hsv, **kwargs)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #21afea61
|
|
class _Hue():
|
|
def __init__(self, max_hue=0.1, p=0.75, draw=None, batch=False): store_attr()
|
|
|
|
def _def_draw(self, x):
|
|
if not self.batch: res = x.new_empty(x.size(0)).uniform_(math.log(1-self.max_hue), -math.log(1-self.max_hue))
|
|
else: res = x.new_zeros(x.size(0)) + random.uniform(math.log(1-self.max_hue), -math.log(1-self.max_hue))
|
|
return torch.exp(res)
|
|
|
|
def before_call(self, x):
|
|
self.change = _draw_mask(x, self._def_draw, draw=self.draw, p=self.p, neutral=0., batch=self.batch)
|
|
|
|
def __call__(self, x):
|
|
h,s,v = x.unbind(1)
|
|
h += self.change[:,None,None]
|
|
h = h % 1.0
|
|
return x.set_(torch.stack((h, s, v),dim=1))
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #ed80ab71
|
|
@patch
|
|
@delegates(_Hue.__init__)
|
|
def hue(x: TensorImage, **kwargs):
|
|
func = _Hue(**kwargs)
|
|
func.before_call(x)
|
|
return TensorImage(x.hsv(func))
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #2e444b4a
|
|
class Hue(HSVTfm):
|
|
"Apply change in hue of `max_hue` to batch of images with probability `p`."
|
|
# Ref: https://pytorch.org/docs/stable/torchvision/transforms.html#torchvision.transforms.functional.adjust_hue
|
|
def __init__(self,
|
|
max_hue:float=0.1, # Maximum scale of changing Hue
|
|
p:float=0.75, # Probability of appying transformation
|
|
draw:float|MutableSequence|Callable=None, # User defined behavior of batch transformation
|
|
batch=False # Apply identical Hue to entire batch
|
|
):
|
|
super().__init__(_Hue(max_hue, p, draw, batch))
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #b27f380b
|
|
def cutout_gaussian(
|
|
x:Tensor, # Input image
|
|
areas:list # List of areas to cutout. Order rl,rh,cl,ch
|
|
):
|
|
"Replace all `areas` in `x` with N(0,1) noise"
|
|
chan,img_h,img_w = x.shape[-3:]
|
|
for rl,rh,cl,ch in areas: x[..., rl:rh, cl:ch].normal_()
|
|
return x
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #d74b7127
|
|
def norm_apply_denorm(
|
|
x:Tensor, # Input Image
|
|
f:Callable, # Function to apply
|
|
nrm:Callable # Normalization transformation
|
|
):
|
|
"Normalize `x` with `nrm`, then apply `f`, then denormalize"
|
|
y = f(nrm(x.clone()))
|
|
return nrm.decode(y).clamp(0,1)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #2e3ba2fb
|
|
def _slice(area, sz):
|
|
bound = int(round(math.sqrt(area)))
|
|
loc = random.randint(0, max(sz-bound, 0))
|
|
return loc,loc+bound
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #38f57374
|
|
class RandomErasing(RandTransform):
|
|
"Randomly selects a rectangle region in an image and randomizes its pixels."
|
|
order = 100 # After Normalize
|
|
def __init__(self,
|
|
p:float=0.5, # Probability of appying Random Erasing
|
|
sl:float=0., # Minimum proportion of erased area
|
|
sh:float=0.3, # Maximum proportion of erased area
|
|
min_aspect:float=0.3, # Minimum aspect ratio of erased area
|
|
max_count:int=1 # Maximum number of erasing blocks per image, area per box is scaled by count
|
|
):
|
|
store_attr()
|
|
super().__init__(p=p)
|
|
self.log_ratio = (math.log(min_aspect), math.log(1/min_aspect))
|
|
|
|
def _bounds(self, area, img_h, img_w):
|
|
r_area = random.uniform(self.sl,self.sh) * area
|
|
aspect = math.exp(random.uniform(*self.log_ratio))
|
|
return _slice(r_area*aspect, img_h) + _slice(r_area/aspect, img_w)
|
|
|
|
def encodes(self,x:TensorImage):
|
|
count = random.randint(1, self.max_count)
|
|
_,img_h,img_w = x.shape[-3:]
|
|
area = img_h*img_w/count
|
|
areas = [self._bounds(area, img_h, img_w) for _ in range(count)]
|
|
return cutout_gaussian(x, areas)
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #df9c77be
|
|
def _compose_same_tfms(tfms):
|
|
tfms = L(tfms)
|
|
if len(tfms) == 0: return None
|
|
res = tfms[0]
|
|
for tfm in tfms[1:]: res.compose(tfm)
|
|
return res
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #947ffcfc
|
|
def setup_aug_tfms(tfms):
|
|
"Go through `tfms` and combines together affine/coord or lighting transforms"
|
|
aff_tfms = [tfm for tfm in tfms if isinstance(tfm, AffineCoordTfm)]
|
|
lig_tfms = [tfm for tfm in tfms if isinstance(tfm, LightingTfm)]
|
|
others = [tfm for tfm in tfms if tfm not in aff_tfms+lig_tfms]
|
|
lig_tfm = _compose_same_tfms(lig_tfms)
|
|
aff_tfm = _compose_same_tfms(aff_tfms)
|
|
res = [aff_tfm] if aff_tfm is not None else []
|
|
if lig_tfm is not None: res.append(lig_tfm)
|
|
return res + others
|
|
|
|
# %% ../../nbs/09_vision.augment.ipynb #ffaa32c1
|
|
def aug_transforms(
|
|
mult:float=1.0, # Multiplication applying to `max_rotate`,`max_lighting`,`max_warp`
|
|
do_flip:bool=True, # Random flipping
|
|
flip_vert:bool=False, # Flip vertically
|
|
max_rotate:float=10., # Maximum degree of rotation
|
|
min_zoom:float=1., # Minimum zoom
|
|
max_zoom:float=1.1, # Maximum zoom
|
|
max_lighting:float=0.2, # Maximum scale of changing brightness
|
|
max_warp:float=0.2, # Maximum value of changing warp per
|
|
p_affine:float=0.75, # Probability of applying affine transformation
|
|
p_lighting:float=0.75, # Probability of changing brightnest and contrast
|
|
xtra_tfms:list=None, # Custom Transformations
|
|
size:int|tuple=None, # Output size, duplicated if one value is specified
|
|
mode:str='bilinear', # PyTorch `F.grid_sample` interpolation
|
|
pad_mode=PadMode.Reflection, # A `PadMode`
|
|
align_corners=True, # PyTorch `F.grid_sample` align_corners
|
|
batch=False, # Apply identical transformation to entire batch
|
|
min_scale=1. # Minimum scale of the crop, in relation to image area
|
|
):
|
|
"Utility func to easily create a list of flip, rotate, zoom, warp, lighting transforms."
|
|
res,tkw = [],dict(size=size if min_scale==1. else None, mode=mode, pad_mode=pad_mode, batch=batch, align_corners=align_corners)
|
|
max_rotate,max_lighting,max_warp = array([max_rotate,max_lighting,max_warp])*mult
|
|
if do_flip: res.append(Dihedral(p=0.5, **tkw) if flip_vert else Flip(p=0.5, **tkw))
|
|
if max_warp: res.append(Warp(magnitude=max_warp, p=p_affine, **tkw))
|
|
if max_rotate: res.append(Rotate(max_deg=max_rotate, p=p_affine, **tkw))
|
|
if min_zoom<1 or max_zoom>1: res.append(Zoom(min_zoom=min_zoom, max_zoom=max_zoom, p=p_affine, **tkw))
|
|
if max_lighting:
|
|
res.append(Brightness(max_lighting=max_lighting, p=p_lighting, batch=batch))
|
|
res.append(Contrast(max_lighting=max_lighting, p=p_lighting, batch=batch))
|
|
if min_scale!=1.: xtra_tfms = RandomResizedCropGPU(size, min_scale=min_scale, ratio=(1,1)) + L(xtra_tfms)
|
|
return setup_aug_tfms(res + L(xtra_tfms))
|