"""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))