<jupyter_start><jupyter_text>Transform refactoring  Setup (See final section of notebook for one-time data processing steps.)<jupyter_code>PATH = Path('data/carvana_small')
PATH_PNG = PATH/'train_masks_png'
PATH_X_FULL = PATH/'train'
PATH_X_128 = PATH/'train-128'
PATH_Y_FULL = PATH_PNG
PATH_Y_128 = PATH/'train_masks-128'

# start with the 128x128 images
PATH_X = PATH_X_128
PATH_Y = PATH_Y_128
class ItemBase():
    @property
    @abstractmethod
    def data(self): pass
    @property
    def tfm_meta(self): return {
        'type': 'base'
    }
    
    def apply_tfms(self, tfms, do_resolve=True, xtra=None, **kwargs):
        if tfms or xtra or size:
            if not xtra: xtra={}
            tfms = sorted(listify(tfms), key=lambda o: o.tfm_order)
            if do_resolve: resolve_tfms(tfms)
            x = self.clone()
            for tfm in tfms:
                if tfm in xtra: x = tfm(x, **xtra[tfm])
                else: x = tfm(x)
        return x
    

class ImageBase(ItemBase):
    @property
    def tfm_meta(self): return {
        'type': 'imagebase',
        'image':True,
        'affine': False,
        'lighting': False,
        'coord': False,
        'pixel': False,
        'data_type':'pixels'
    }
    def lighting(self, func, *args, **kwargs): return self
    def pixel(self, func, *args, **kwargs): return self
    def coord(self, func, *args, **kwargs): return self
    def affine(self, func, *args, **kwargs): return self

    def set_sample(self, **kwargs):
        self.sample_kwargs = kwargs
        return self
    
    def clone(self): return self.__class__(self.data.clone())
    
    def apply_tfms(self, tfms, do_resolve=True, xtra=None, size=None,
                   mult=32, do_crop=True, padding_mode='reflect', **kwargs):
        if tfms or xtra or size:
            if not xtra: xtra={}
            tfms = sorted(listify(tfms), key=lambda o: o.tfm_order)
            if do_resolve: resolve_tfms(tfms)
            x = self.clone()
            x.set_sample(padding_mode=padding_mode, **kwargs)
            if size:
                crop_target = get_crop_target(size, mult=mult)
                target = get_resize_target(x, crop_target, do_crop=do_crop)
                x.resize(target)
            size_tfms = [o for o in tfms if o.tfm_type == 'crop']
            for tfm in tfms:
                if tfm in xtra: x = tfm(x, **xtra[tfm])
                elif tfm in size_tfms: x = tfm(x, size=size, padding_mode=padding_mode)
                else: x = tfm(x)
        return x
            
            
    
class Image(ImageBase):
    @property
    def tfm_meta(self): return {
        'type': 'image',
        'image':True,
        'affine': True,
        'lighting': True,
        'coord': True,
        'pixel': True,
        'data_type':'pixels'
    }

    def __init__(self, px):
        self._px = px
        self._logit_px=None
        self._flow=None
        self._affine_mat=None
        self.sample_kwargs = {}

    @property
    def shape(self): return self._px.shape
    
    def __repr__(self): return f'{self.__class__.__name__} ({self.shape})'

    def refresh(self):
        if self._logit_px is not None:
            self._px = self._logit_px.sigmoid_()
            self._logit_px = None
        if self._affine_mat is not None or self._flow is not None:
            self._px = grid_sample(self._px, self.flow, **self.sample_kwargs)
            self.sample_kwargs = {}
            self._flow = None
        return self

    @property
    def px(self):
        self.refresh()
        return self._px
    @px.setter
    def px(self,v): self._px=v

    @property
    def flow(self):
        if self._flow is None:
            self._flow = affine_grid(self.shape)
        if self._affine_mat is not None:
            self._flow = affine_mult(self._flow,self._affine_mat)
            self._affine_mat = None
        return self._flow
    @flow.setter
    def flow(self,v): self._flow=v

    def lighting(self, func, *args, **kwargs):
        self.logit_px = func(self.logit_px, *args, **kwargs)
        return self

    def pixel(self, func, *args, **kwargs):
        self.px = func(self.px, *args, **kwargs)
        return self

    def coord(self, func, *args, **kwargs):
        self.flow = func(self.flow, self.shape, *args, **kwargs)
        return self

    def affine(self, func, *args, **kwargs):
        m = func(*args, **kwargs)
        self.affine_mat = self.affine_mat @ self._px.new(m)
        return self

    def resize(self, size):
        assert self._flow is None
        if isinstance(size, int): size=(self.shape[0], size, size)
        self.flow = affine_grid(size)
        return self

    @property
    def affine_mat(self):
        if self._affine_mat is None: self._affine_mat = self._px.new(torch.eye(3))
        return self._affine_mat
    @affine_mat.setter
    def affine_mat(self,v): self._affine_mat=v

    @property
    def logit_px(self):
        if self._logit_px is None: self._logit_px = logit_(self.px)
        return self._logit_px
    @logit_px.setter
    def logit_px(self,v): self._logit_px=v
    
    def show(self, ax=None, **kwargs): show_image(self.px, ax=ax, **kwargs)
    
    @property
    def data(self): return self.px
    
class ImageMask(Image):
    @property
    def tfm_meta(self): return {
        'type': 'mask',
        'image':True,
        'affine': True,
        'lighting': False,
        'coord': True,
        'pixel': True,
        'data_type':'mask'
    }
    
    def lighting(self, func, *args, **kwargs): 
        assert('never call')
        return self
    
    def refresh(self):
        self.sample_kwargs['mode'] = 'nearest'
        return super().refresh()

def open_mask(fn):
    return ImageMask(pil2tensor(PIL.Image.open(fn)).float())

import nb_002
nb_002.Image = Image
img_f = next(PATH_X.iterdir())
x = open_image(img_f)
x.show()
x.shape
def get_y_fn(x_fn): return PATH_Y/f'{x_fn.name[:-4]}_mask.png'

img_y_f = get_y_fn(img_f)
y = open_mask(img_y_f)
y.show()
#export
# Same as `show_image`, but renamed with _ prefix
def _show_image(img, ax=None, figsize=(3,3), hide_axis=True, cmap='binary', alpha=None):
    if ax is None: fig,ax = plt.subplots(figsize=figsize)
    ax.imshow(image2np(img), cmap=cmap, alpha=alpha)
    if hide_axis: ax.axis('off')
    return ax

def show_image(x, y=None, ax=None, figsize=(3,3), alpha=0.5, hide_axis=True, cmap='viridis'):
    ax1 = _show_image(x, ax=ax, hide_axis=hide_axis, cmap=cmap)
    if y is not None: _show_image(y, ax=ax1, alpha=alpha, hide_axis=hide_axis, cmap=cmap)
    if hide_axis: ax1.axis('off')
        
def _show(self, ax=None, y=None, **kwargs):
    if y is not None: y=y.data
    return show_image(self.data, ax=ax, y=y, **kwargs)

Image.show = _show
x.show(y=y)<jupyter_output><empty_output><jupyter_text>Type disptach<jupyter_code>#export
from functools import wraps

def get_default_args(func):
    if hasattr(func, '_def_args'): return func._def_args
    else: return {k: v.default
            for k, v in inspect.signature(func).parameters.items()
            if v.default is not inspect.Parameter.empty}

def get_arg_names(func):
    return list(inspect.signature(func).parameters)

def get_dynamic_var_args(func):
    return {k: v.default
            for k, v in inspect.signature(func).parameters.items()}

def bind_args(func, v):
    arg_names = get_arg_names(func)
    bound_args = { arg_names[i]: vi for i, vi in enumerate(v)}
    return bound_args
        

def resolve_fun_annotations(func, kwargs):
    params = copy(func.__annotations__)
    resolved = {}
    for k, v in kwargs.items():
        if k in params and not isinstance(v, Dynamic):
            rand_func = params[k]
            if isinstance(v, Dict): resolved[k] = rand_func(**v)
            elif isinstance(v, Iterable):
                arg_names = get_arg_names(rand_func)
                bound_args = bind_args(rand_func, v)
                resolved[k] = rand_func(**bound_args)
            elif rand_func == Bool and type(v) != bool: 
                resolved[k] = rand_func(v)
            else: resolved[k] = v
        else: resolved[k] = v
    return resolved


def dynamic_resolve(a):
    if isinstance(a, Dynamic): return a.resolve()
    else: return a

def dynamic_release(a):
    if isinstance(a, Dynamic): return a.release()
    elif isinstance(a, list):
        for ai in a: dynamic_release(ai)
    elif isinstance(a, dict):
        for vi in a.values(): dynamic_release(vi)
            
class Dynamic(object):
    def __init__(self, func=None, args=None, kwargs=None):
        self.func = func       
        self.args = args
        self.kwargs = resolve_fun_annotations(self.func, kwargs) if kwargs else {}
        self.def_args = get_default_args(self.func) 
        self.init_state()
        
    def init_state(self):
        self.value_ = None
        self.bound = False
        self.bound_args = {}
        self.bound_kwargs = {}

    def clone(self, **kwargs):
        copy_kwargs = {k:copy(v) for k,v in self.kwargs.items()}
        copy_args = copy(self.args)
        kwargs = {**copy_kwargs, **kwargs}
        dtype = type(self)
        new_dynamic = dtype(func=self.func, args=copy_args, kwargs=kwargs)
        new_dynamic.__dict__.update({k:v for k,v in self.__dict__.items()
                                     if k not in new_dynamic.__dict__})
        return new_dynamic

    def override(self, **kwargs):
        kwargs = {**self.kwargs, **kwargs}
        dtype = type(self)
        new_dynamic = dtype(func=self.func, args=self.args, kwargs=kwargs)
        new_dynamic.__dict__.update({k:v for k,v in self.__dict__.items()
                                     if k not in new_dynamic.__dict__})
        return new_dynamic

    def __copy__(self):
        return self.clone()
    
    def needs_bind(self, overrides=None):
        return not self.bound
    
    def bind(self, **kwargs):
        if self.kwargs:
            kwargs = resolve_fun_annotations(self.func, kwargs) 
            bound_overrides = {k:dynamic_resolve(v)for k,v in kwargs.items()}
        else: bound_overrides = {}
        if self.needs_bind(bound_overrides):
            self.bound_args = bind_args(self.func, self.args) if self.args else {}
            bind_kw = {**self.def_args, **self.kwargs, **self.bound_args}
            self.bound_kwargs = {k:dynamic_resolve(v)for k,v in bind_kw.items()}
            self.value = None
        self.bound = True
        return bound_overrides
     
    
    def release(self):
        dynamic_release(self.kwargs)
        dynamic_release(self.bound_args)
        dynamic_release(self.def_args)
        dynamic_release(self.func)
        self.init_state()
    
    def __repr__(self):
        return f'{self.func.__name__}:{self.kwargs}:{self.value_}'
    
                
class DynamicVar(Dynamic):      
    def __init__(self, func=None, args=None, kwargs=None):
        super().__init__(func=func, args=args, kwargs=kwargs)
        
    def resolve(self):
        self.bind()
        if self.value_ is None:
            self.value_ = self.func(**self.bound_kwargs)
        return self.value_


class DynamicFunc(Dynamic):      
    def __init__(self, func=None, args=None, kwargs=None):
        assert(args is None)
        super().__init__(func=func, args=args, kwargs=kwargs)
 
    def needs_bind(self, overrides=None):
        if not self.bound: return True
        for v in self.bound_args.values():
            if isinstance(v, Dynamic) and v.needs_bind(): return True
        for v in self.bound_kwargs.values():
            if isinstance(v, Dynamic) and v.needs_bind(): return True
        if overrides:
            for v in overrides.values():
                if isinstance(v, Dynamic) and v.needs_bind(): return True
        return False

    def resolve(self):
        self.release()
        self.bind()
        return self
    
    def __call__(self, *args, **kwargs):
        if args:
            # force functions to always recheck kwargs
            # could be released underneath
            bind_overrides = self.bind(**kwargs)
            bound_kw = {**self.bound_kwargs, **bind_overrides}
            return self.func(*args, **bound_kw)
        else:
            new_func = self.override(**kwargs)
            return new_func


def dynamic_var(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        return DynamicVar(func=func, args=args, kwargs=kwargs) 
    return wrapper

def dynamic_fn(func):
    return DynamicFunc(func=func)

def uniform(low, high, size=None):
    return random.uniform(low,high) if size is None else torch.FloatTensor(size).uniform_(low,high)

def uniform_int(low, high, size=None):
    return random.randint(low,high) if size is None else torch.randint(low,high,size)

@dynamic_var
def Uniform(low, high, size=None):
    return uniform(low, high, size=size)

@dynamic_var
def UniformInt(low, high, size=None):
    return uniform_int(low, high, size=size)

@dynamic_var
def RandInt(low, high):
    return random.randint(low, high)
    

@dynamic_var
def LogUniform(low, high, size=None):
    res = uniform(log(low), log(high), size)
    value = exp(res) if size is None else res.exp_()
    return value

@dynamic_var
def Bool(p, size=None):
    return uniform(0, 1, size=size)<p

@dynamic_var
def Const(c):
    return c

@dynamic_var
def Choice(choices):
    return random.choice(choices)


@dynamic_fn
def maybe(x, *args, f:DynamicFunc=None, p:Bool=True, **kwargs):
    if p: return f(x, *args, **kwargs)
    else: return x

def make_rand(f, *args, p:Bool=True, **kwargs):
    f = f(**kwargs)
    new_fn = maybe(f=f, p=p)
    if args: return new_fn(*args)
    else: return new_fn
class Tfm(object):
    def __init__(self, func, dispatch, attrs={}):
        self.default_func = func
        self.dispatch = dispatch
        self._reg_funcs = {}
        self._attrs=attrs
    
    def reg(self, disp_id):
        def _reg_type(func):
            self._reg_funcs[disp_id] = func
        return _reg_type
    
    def wrap_call(self, fn, *args, **kwargs):
        return fn(*args, **kwargs)
        
    def __call__(self, *args, p:Bool=True, **kwargs):
        if args:
            fn_key = self.dispatch(*args, **kwargs)
            fn = self._reg_funcs.get(fn_key, self.default_func)
            fn = DynamicFunc(fn,kwargs=kwargs)
            fn = maybe(f=fn, p=p)
            return self.wrap_call(fn, *args)
        else: 
            @dynamic_fn
            @wraps(self.default_func)
            def _dyn_tfm(*t_args, **t_kwargs):
                t_kwargs = { **kwargs, **t_kwargs}
                fn_key = self.dispatch(*t_args, **t_kwargs)
                fn = self._reg_funcs.get(fn_key, self.default_func)
                return self.wrap_call(fn, *t_args, **t_kwargs)              
            fn = _dyn_tfm(**kwargs)
            fn = maybe(f=fn, p=p)
            for k,v in self._attrs.items(): 
                setattr(fn, k, v)
            return fn
          
                
class TfmImage(Tfm):
    def __init__(self, func, dispatch, wrap=None, order=None, tfm_type=None, **attrs):
        self.wrap = wrap
        super().__init__(func, dispatch, {**attrs, 'tfm_order': order, 'tfm_type': tfm_type})
            
    def wrap_call(self, fn, *args, **kwargs):
        if self.wrap:
            # check if this img obj supports this wrap type
            img = args[0]
            if img.tfm_meta.get(self.wrap, False):
                wrap_fn = getattr(args[0], self.wrap)
                if wrap_fn: 
                    wrap_fn(fn, *args[1:], **kwargs)
                    return img
            # wrap unsupported just return img
            else: return img
        return fn(*args, **kwargs)

def releasee_tfms(tfms):
    for f in listify(tfms): f.release()



def disp_on_type(x, *args, **kwargs): return type(x)
TypeTfm = partial(Tfm, dispatch=disp_on_type)

def disp_on_img_meta(x, *args, **kwargs):
    if isinstance(x, ImageBase): return x.tfm_meta['data_type']
    else: return None

def make_img_tfm(tfm_type=None, **kwargs):
    if not tfm_type: tfm_type = kwargs.get('wrap', 'unknown')
    return partial(TfmImage, dispatch=disp_on_img_meta, **{**kwargs, 'tfm_type': tfm_type})

TfmLighting = make_img_tfm(wrap='lighting', order=9)
TfmAffine = make_img_tfm(wrap='affine', order=5)
TfmPixel = make_img_tfm(wrap='pixel', order=10)
TfmPixelPad = make_img_tfm(wrap='pixel', order=-10)
TfmCoord = make_img_tfm(wrap='coord', order=4)
TfmPad = make_img_tfm(wrap='pixel', order=-10)
TfmCrop = make_img_tfm(wrap='pixel', order=99, tfm_type='crop')


@TypeTfm
def mult(x, a:Uniform=2):
    print('mult', x, a)
    return a*x

@mult.reg(str)
def mult_str(x, a):
    print('string', x, a)
    return x * int(abs(a))

def _rotate(degrees, forward=True):
    flip = 1. if forward else -1.
    angle = degrees * math.pi / 180
    return [[cos(angle),       flip*-sin(angle), 0.],
            [flip*sin(angle),  cos(angle),       0.],
            [0.,               0.        ,       1.]]

@TfmAffine
def rotate(degrees:Uniform=0., forward=True): 
    return _rotate(degrees, forward)

@TfmLighting
def brightness(x, change:Uniform=0., forward=True): 
    return x.add_(scipy.special.logit(change))

@TfmLighting
def contrast(x, scale:LogUniform=1., forward=True): return x.mul_(scale)

def get_zoom_mat(sw, sh, c, r, forward=True):
    if not forward:
        sw = 1/sw
        sh = 1/sh
    return [[sw, 0,  c],
            [0, sh,  r],
            [0,  0, 1.]]

@TfmAffine
def zoom(x, scale:Uniform=1.0, 
         row_pct:Uniform=0.5, col_pct:Uniform=0.5, forward=True):
    s = 1-1/scale
    col_c = s * (2*col_pct - 1)
    row_c = s * (2*row_pct - 1)
    return get_zoom_mat(1/scale, 1/scale, col_c, row_c)

@TfmAffine
def squish(x, scale:Uniform=1.0, 
           row_pct:Uniform=0.5, col_pct:Uniform=0.5, forward=True):
    if scale <= 1: 
        col_c = (1-scale) * (2*col_pct - 1)
        return get_zoom_mat(scale, 1, col_c, 0.)
    else:          
        row_c = (1-1/scale) * (2*row_pct - 1)
        return get_zoom_mat(1, 1/scale, 0., row_c)
    
@TfmCoord
def jitter(x, magnitude:Uniform, forward=True):
    return c.add_((torch.rand_like(c)-0.5)*magnitude*2)

@TfmPixel
def flip_lr(x, forward=True): 
    return x.flip(2)

@TfmPixelPad
def pad(x, padding=0, mode='reflect'): 
    return F.pad(x[None], (padding,)*4, mode=mode)[0]

@TfmPixel
def crop(x, size, row_pct:Uniform=0.5, col_pct:uniform=0.5): 
    size = listify(size,2)
    rows,cols = size
    row = int((x.size(1)-rows+1) * row_pct)
    col = int((x.size(2)-cols+1) * col_pct)
    return x[:, row:row+rows, col:col+cols].contiguous()

def compute_zs_mat(sz, scale, squish, invert, row_pct, col_pct):
    orig_ratio = math.sqrt(sz[2]/sz[1])
    for s,r,i in zip(scale,squish, invert):
        s,r = math.sqrt(s),math.sqrt(r)
        if s * r <= 1 and s / r <= 1: #Test if we are completely inside the picture
            w,h = (s/r, s*r) if i else (s*r,s/r)
            w /= orig_ratio
            h *= orig_ratio
            col_c = (1-w) * (2*col_pct - 1)
            row_c = (1-h) * (2*row_pct - 1)
            return get_zoom_mat(w, h, col_c, row_c)
        
    #Fallback, hack to emulate a center crop without cropping anything yet.
    if orig_ratio > 1: return get_zoom_mat(1/orig_ratio**2, 1, 0, 0.)
    else:              return get_zoom_mat(1, orig_ratio**2, 0, 0.)

@TfmCoord
def zoom_squish(x, scale:Uniform=1.0, squish:Uniform=1.0, invert:Bool=False, 
                row_pct:Uniform=0.5, col_pct:Uniform=0.5):
    #This is intended for scale, squish and invert to be of size 10 (or whatever) so that the transform
    #can try a few zoom/squishes before falling back to center crop (like torchvision.RandomResizedCrop)
    m = compute_zs_mat(size, scale, squish, invert, row_pct, col_pct)
    return affine_mult(c, FloatTensor(m))

_orig_pts = [[-1,-1], [-1,1], [1,-1], [1,1]]

def _perspective_warp(c, targ_pts):
    return apply_perspective(c, find_coeffs(_orig_pts, targ_pts))

@TfmCoord
def perspective_warp(x, magnitude:partial(Uniform,size=8)=0):
    magnitude = magnitude.view(4,2)
    targ_pts = [[x+m for x,m in zip(xs, ms)] for xs, ms in zip(_orig_pts, magnitude)]
    return _perspective_warp(c, targ_pts)

@TfmCoord
def symmetric_warp(x, magnitude:partial(Uniform,size=4)=0):
    m = listify(magnitude, 4)
    targ_pts = [[-1-m[3],-1-m[1]], [-1-m[2],1+m[1]], [1+m[3],-1-m[0]], [1+m[2],1+m[0]]]
    return _perspective_warp(c, targ_pts)

@TfmCoord
def tilt(c, img_size, direction:rand_int, magnitude:Uniform=0):
    orig_pts = [[-1,-1], [-1,1], [1,-1], [1,1]]
    if direction == 0:   targ_pts = [[-1,-1], [-1,1], [1,-1-magnitude], [1,1+magnitude]]
    elif direction == 1: targ_pts = [[-1,-1-magnitude], [-1,1+magnitude], [1,-1], [1,1]]
    elif direction == 2: targ_pts = [[-1,-1], [-1-magnitude,1], [1,-1], [1+magnitude,1]]
    elif direction == 3: targ_pts = [[-1-magnitude,-1], [-1,1], [1+magnitude,-1], [1,1]]  
    coeffs = find_coeffs(orig_pts, targ_pts)
    return apply_perspective(c, coeffs)

@TfmCoord
def skew(c, img_size, direction:RandInt, magnitude:Uniform=0):
    orig_pts = [[-1,-1], [-1,1], [1,-1], [1,1]]
    if direction == 0:   targ_pts = [[-1-magnitude,-1], [-1,1], [1,-1], [1,1]]
    elif direction == 1: targ_pts = [[-1,-1-magnitude], [-1,1], [1,-1], [1,1]]
    elif direction == 2: targ_pts = [[-1,-1], [-1-magnitude,1], [1,-1], [1,1]]
    elif direction == 3: targ_pts = [[-1,-1], [-1,1+magnitude], [1,-1], [1,1]]
    elif direction == 4: targ_pts = [[-1,-1], [-1,1], [1+magnitude,-1], [1,1]]
    elif direction == 5: targ_pts = [[-1,-1], [-1,1], [1,-1-magnitude], [1,1]]
    elif direction == 6: targ_pts = [[-1,-1], [-1,1], [1,-1], [1+magnitude,1]]
    elif direction == 7: targ_pts = [[-1,-1], [-1,1], [1,-1], [1,1+magnitude]] 
    coeffs = find_coeffs(orig_pts, targ_pts)
    return apply_perspective(c, coeffs)
    
@TfmCrop
def crop_pad(x, size=None, padding_mode='reflect',
             row_pct:Uniform = 0.5, col_pct:Uniform = 0.5):
    if padding_mode=='zeros': padding_mode='constant'
    size = listify(size,2)
    if x.shape[1:] == size: return x
    rows,cols = size
    if x.size(1)<rows or x.size(2)<cols:
        row_pad = max((rows-x.size(1)+1)//2, 0)
        col_pad = max((cols-x.size(2)+1)//2, 0)
        x = F.pad(x[None], (col_pad,col_pad,row_pad,row_pad), mode=padding_mode)[0]
    row = int((x.size(1)-rows+1)*row_pct)
    col = int((x.size(2)-cols+1)*col_pct)

    x = x[:, row:row+rows, col:col+cols]
    return x.contiguous() # without this, get NaN later - don't know why


@TfmPixel
def dihedral(x, k:partial(uniform_int,0,8)):
    flips=[]
    if k&1: flips.append(1)
    if k&2: flips.append(2)
    if flips: x = torch.flip(x,flips)
    if k&4: x = x.transpose(1,2)
    return x.contiguous()

def rand_zoom(*args, **kwargs): return zoom(*args, row_pct=(0,1), col_pct=(0,1), **kwargs)
def rand_crop(*args, **kwargs): return crop_pad(*args, row_pct=(0,1), col_pct=(0,1), **kwargs)
def zoom_crop(scale, do_rand=False, p=1.0):
    zoom_fn = rand_zoom if do_rand else zoom
    crop_fn = rand_crop if do_rand else crop_pad
    return [zoom_fn(scale=scale, p=p), crop_fn()]

import nb_002, nb_002c, nb_005, nb_003


nb_002.rotate = rotate
nb_002.brightness = brightness
nb_002.contrast = contrast
nb_002.zoom = zoom
nb_002.squish = squish
nb_002c.skew = skew
nb_002c.tilt = tilt
nb_002c.symmetric_warp = symmetric_warp
nb_002c.perspective_warp = perspective_warp
nb_003.crop = crop
nb_003.crop_pad = crop_pad
nb_003.pad = pad
nb_003.rand_crop = rand_crop
nb_003.rand_zoom = rand_zoom
nb_003.zoom_crop = zoom_crop
nb_005.dihedral = dihedral

brightness(x.clone(), change=(0.8,0.99)).show()
b = brightness(change=(0.8,0.99))
print(b.tfm_order, b.tfm_type)
b(x.clone()).set_sample(padding_mode='zeros').show()
b(y.clone()).set_sample(padding_mode='zeros').show()
r = rotate(degrees=(-50,50))
r(x.clone()).set_sample(padding_mode='zeros').show()
r(y.clone()).set_sample(padding_mode='zeros').show()
def xi(): return x.clone()

tfm = rotate(degrees=(-30,30), p=0.5)
_,axes = plt.subplots(2,4, figsize=(12,6))
for i,ax in enumerate(axes.flat): 
    tfm(xi()).set_sample(padding_mode='zeros').show(ax)
    tfm.release()
def xi(): return x.clone()
tfm = crop_pad(size=200, row_pct=0.,col_pct=0., padding_mode='constant')
#tfm = dihedral()
tfm(xi()).show()
_,axes = plt.subplots(4,4, figsize=(12,6))
for i,ax in enumerate(axes.flatten()):
    magns = [-0.4,0.4]
    skew(xi(),
         direction=i%8, 
         magnitude=magns[i//8]
    ).set_sample(padding_mode='zeros').show(ax)
_,axes = plt.subplots(4,4, figsize=(12,6))
for i,ax in enumerate(axes.flatten()):
    tfm = skew(
        direction=i%8, 
        magnitude=magns[i//8]
    )
    tfm(xi()).set_sample(padding_mode='zeros').show(ax)
def get_transforms(do_flip=False, flip_vert=False, max_rotate=0., max_zoom=1., max_lighting=0., max_warp=0.,
                   p_affine=0.75, p_lighting=0.5, xtra_tfms=None):
    res = [rand_crop()]
    if do_flip:    res.append(dihedral() if flip_vert else flip_lr(p=0.5))
    if max_warp:   res.append(symmetric_warp(magnitude=(-max_warp,max_warp), p=p_affine))
    if max_rotate: res.append(rotate(degrees=(-max_rotate,max_rotate), p=p_affine))
    if max_zoom>1: res.append(rand_zoom(scale=(1.,max_zoom), p=p_affine))
    if max_lighting:
        res.append(brightness(change=(0.5*(1-max_lighting), 0.5*(1+max_lighting)), p=p_lighting))
        res.append(contrast(scale=(1-max_lighting, 1/(1-max_lighting)), p=p_lighting))
    #       train                   , valid
    return (res + listify(xtra_tfms), [crop_pad()])
get_transforms

import nb_005
nb_005.get_transforms = get_transforms<jupyter_output><empty_output><jupyter_text>Dataset - data types: regr, class, seg, bbox, polygon, generative (s/res, color), custom<jupyter_code>#export
class DatasetTfm(Dataset):
    def __init__(self, ds:Dataset, tfms:Collection[Callable]=None, tfm_y:bool=False, **kwargs):
        self.ds,self.tfms,self.tfm_y,self.x_kwargs = ds,tfms,tfm_y,kwargs
        self.y_kwargs = {**self.x_kwargs, 'do_resolve':False} # don't reset random vars
        
    def __len__(self): return len(self.ds)
    
    def __getitem__(self,idx):
        x,y = self.ds[idx]
        x = x.apply_tfms(self.tfms, **self.x_kwargs)
        if self.tfm_y: y = y.apply_tfms(self.tfms, **self.y_kwargs)
        return x.data, y.data
    
    def __getattr__(self,k): return getattr(self.ds, k)
    
import nb_002b,nb_005
nb_002b.DatasetTfm = DatasetTfm  
nb_005.DatasetTfm  = DatasetTfm  
#export
class MatchedImageDataset(DatasetBase):
    def __init__(self, x:Collection[Path], y:Collection[Path]):
        assert len(x)==len(y)
        self.x,self.y = np.array(x),np.array(y)

    def __getitem__(self, i): 
        return open_image(self.x[i]), open_mask(self.y[i])
    
def split_arrs(idxs, *a):
    mask = np.zeros(len(a[0]),dtype=bool)
    mask[np.array(idxs)] = True
    return [(o[mask],o[~mask]) for o in map(np.array, a)]
def get_datasets(path):
    x_fns = [o for o in path.iterdir() if o.is_file()]
    y_fns = [get_y_fn(o) for o in x_fns]
    val_idxs = list(range(1008))
    ((val_x,trn_x),(val_y,trn_y)) = split_arrs(val_idxs, x_fns, y_fns)
    return (MatchedImageDataset(trn_x, trn_y),
            MatchedImageDataset(val_x, val_y))
train_ds,valid_ds = get_datasets(PATH_X_128)
train_ds,valid_ds
x,y = next(iter(train_ds))
x.shape, y.shape, type(x), type(y)
size=128
def get_tfm_datasets(size):
    datasets = get_datasets(PATH_X_128 if size<=128 else PATH_X_FULL)
    tfms = get_transforms(do_flip=True, max_rotate=4, max_lighting=0.2)
    return transform_datasets(train_ds, valid_ds, tfms, tfm_y=True, size=size, padding_mode='zeros')
train_tds,*_ = get_tfm_datasets(size)
_,axes = plt.subplots(1,4, figsize=(12,6))
for i, ax in enumerate(axes.flat):
    imgx,imgy = train_tds[i]
    Image(imgx).show(ax, y=ImageMask(imgy))
imgx,imgy = train_tds[i]
train_ds, valid_ds = get_datasets(PATH_X_128 if size<=128 else PATH_X_FULL)
imgx,imgy = train_ds[1]
Image(imgx.px.flip(2)).show()
#export
def normalize_batch(b, mean, std, do_y=False):
    x,y = b
    x = normalize(x,mean,std)
    if do_y: y = normalize(y,mean,std)
    return x,y

def normalize_funcs(mean, std, do_y=False, device=None):
    if device is None: device=default_device
    return (partial(normalize_batch, mean=mean.to(device),std=std.to(device), do_y=do_y),
            partial(denormalize,     mean=mean,           std=std))
default_norm,default_denorm = normalize_funcs(*imagenet_stats)
bs = 64
def get_data(size, bs):
    return DataBunch.create(*get_tfm_datasets(size), bs=bs, tfms=default_norm)
data = get_data(size, bs)
#export
def show_xy_images(x,y,rows,figsize=(9,9)):
    fig, axs = plt.subplots(rows,rows,figsize=figsize)
    for i, ax in enumerate(axs.flatten()): show_image(x[i], y=y[i], ax=ax)
    plt.tight_layout()
x,y = next(iter(data.train_dl))
x,y = x.cpu(),y.cpu()
x = default_denorm(x)
show_xy_images(x,y,4, figsize=(9,9))
x.shape, y.shape<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14526817146407123, 'fraction_numerical': 0.018878994960999516, 'mean_word_length': 2.9107841814009987, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '274354', 'license': 'Apache-2.0', 'n_tokens_mistral': 10930, 'n_tokens_neox': 10622, 'n_words': 2292, 'path': 'deep-learning/fastai-docs/fastai_docs-master/dev_nb/experiments/tfm_refactor.ipynb', 'repo_name': 'AadityaGupta/Artificial-Intelligence-Deep-Learning-Machine-Learning-Tutorials', 'size': '28974'}