# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/vision.augment.batch.ipynb.

# %% ../../../nbs/vision.augment.batch.ipynb 1
# Contains code from:
# fastai - Apache License 2.0 - Copyright (c) 2023 fast.ai

# %% ../../../nbs/vision.augment.batch.ipynb 3
from __future__ import annotations

import random
from torch.distributions.uniform import Uniform

from fastai.vision.augment import cutout_gaussian, _slice
from fastai.vision.core import TensorImage

from ...transform import *
from ...imports import *

# %% auto 0
__all__ = ['GrayScaleMode', 'GrayScale', 'ChannelDrop', 'RandomNoise', 'RandomErasingBatch']

# %% ../../../nbs/vision.augment.batch.ipynb 6
class GrayScaleMode(Enum):
    "GrayScaleModes for GrayScale"
    Luma601 = 0
    Luma709 = 1
    Average = 2
    Random  = 3

# %% ../../../nbs/vision.augment.batch.ipynb 7
class GrayScale(BatchRandTransform):
    "Convert RGB image into grayscale using luma_bt.601, luma_bt.709, averaging, or randomly selected"
    order = 55 # After LightingTfms
    def __init__(self,
        p:float=0.1, # Per-item probability
        mode:GrayScaleMode=GrayScaleMode.Random # GrayScaleMode to apply to images. Random applies all three element-wise with equal probability
    ):
        super().__init__(p=p)
        self.mode = mode
        self._luma601 = torch.tensor([0.2989, 0.5870, 0.1140])
        self._luma709 = torch.tensor([0.2126, 0.7152, 0.0722])
        self._average = torch.tensor([0.3333, 0.3333, 0.3333])

    def encodes(self, x:TensorImage):
        if self.mode==GrayScaleMode.Random:
            choice = torch.stack(random.choices([self._luma601, self._luma709, self._average], k=x.shape[0]), dim=0).view(x.shape[0], -1, 1, 1)
            return (x*choice).sum(1)[:,None]
        elif self.mode==GrayScaleMode.Luma601:
            return (x*self._luma601[...,None,None]).sum(1)[:,None]
        elif self.mode==GrayScaleMode.Luma709:
            return (x*self._luma709[...,None,None]).sum(1)[:,None]
        elif self.mode==GrayScaleMode.Average:
            return x.mean(1)[:,None]

    def to(self, *args, **kwargs):
        device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs)
        self._luma601 = self._luma601.to(device=device)
        self._luma709 = self._luma709.to(device=device)
        self._average = self._average.to(device=device)

# %% ../../../nbs/vision.augment.batch.ipynb 9
class ChannelDrop(BatchRandTransform):
    "Drop entire channel by replacing it with random solid value [0,1)"
    order = 65
    def __init__(self,
        p:float=0.1, # Per-item probability
        replace:float|None=None # Set constant replacement value. Defaults to element-wise random value [0,1)
    ):
        super().__init__(p=p)
        self.replace, self._view = replace, None

    def encodes(self, x:TensorImage):
        shape = x.shape
        if self._view is None: self._view = [-1] + [1]*(len(shape)-2)
        value = torch.rand(shape[0], device=x.device)
        idxs = torch.randint(0, shape[1], (shape[0],), device=x.device)
        x[torch.arange(shape[0], device=x.device), idxs, ...] = value.view(self._view)
        return x

# %% ../../../nbs/vision.augment.batch.ipynb 11
class RandomNoise(BatchRandTransform):
    "Add random guassian noise based on `stdev`"
    order = 99 # After Normalize
    def __init__(self,
        p:float=0.25, # Per-item probability
        stdev:float|tuple=(0.1, 0.25), # Maximum or range of the standard deviation of added noise
        random:bool=True # Randomize standard deviation of added noise between [`stdev[0]`, `stdev[1]`)
    ):
        super().__init__(p=p)
        self.stdev, self.random, self.dist = stdev, random, None

    def encodes(self, x:TensorImage):
        shape = x.shape
        if self.random:
            if self.dist is None:
                self.dist = Uniform(torch.tensor(self.stdev[0], device=x.device), torch.tensor(self.stdev[1], device=x.device))
                self.view = [-1, *[1]*len(shape[1:])]
            std = self.dist.sample((shape[0],)).view(self.view)
        else:
            std = self._stdev = self.stdev
        return x + torch.normal(mean=torch.zeros(shape, device=x.device), std=std)

# %% ../../../nbs/vision.augment.batch.ipynb 17
class RandomErasingBatch(BatchRandTransform):
    "Randomly selects a rectangle region in an image and randomizes its pixels."
    order = 100 # After Normalize
    def __init__(self,
        p:float=0.25, # Per-item probability
        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
        element:bool=False # Loop over through batch and apply element-wise unique erasing
    ):
        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):
        _,img_h,img_w = x.shape[-3:]
        if self.element:
            for i in range(x.shape[0]):
                count = random.randint(1, self.max_count)
                area = img_h*img_w/count
                areas = [self._bounds(area, img_h, img_w) for _ in range(count)]
                x[i,...] = cutout_gaussian(x[i,...], areas)
            return x
        else:
            count = random.randint(1, self.max_count)
            area = img_h*img_w/count
            areas = [self._bounds(area, img_h, img_w) for _ in range(count)]
            return cutout_gaussian(x, areas)
