# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/03_vision_augment.ipynb.

# %% auto 0
__all__ = ['CustomDictTransform', 'do_pad_or_crop', 'PadOrCrop', 'ZNormalization', 'BraTSMaskConverter', 'BinaryConverter',
           'RandomGhosting', 'RandomSpike', 'RandomNoise', 'RandomBiasField', 'RandomBlur', 'RandomGamma',
           'RandomMotion', 'RandomElasticDeformation', 'RandomAffine', 'RandomFlip', 'OneOf']

# %% ../nbs/03_vision_augment.ipynb 2
from fastai.data.all import *
from .vision_core import *
import torchio as tio

# %% ../nbs/03_vision_augment.ipynb 5
class CustomDictTransform(ItemTransform):
    """A class that serves as a wrapper to perform an identical transformation on both 
    the image and the target (if it's a mask).
    """
    
    split_idx = 0  # Only perform transformations on training data. Use TTA() for transformations on validation data.

    def __init__(self, aug):
        """Constructs CustomDictTransform object.

        Args:
            aug (Callable): Function to apply augmentation on the image.
        """
        self.aug = aug

    def encodes(self, x):
        """
        Applies the stored transformation to an image, and the same random transformation 
        to the target if it is a mask. If the target is not a mask, it returns the target as is.

        Args:
            x (Tuple[MedImage, Union[MedMask, TensorCategory]]): A tuple containing the 
            image and the target.

        Returns:
            Tuple[MedImage, Union[MedMask, TensorCategory]]: The transformed image and target. 
            If the target is a mask, it's transformed identically to the image. If the target 
            is not a mask, the original target is returned.
        """
        img, y_true = x

        if isinstance(y_true, (MedMask)):
            aug = self.aug(tio.Subject(img=tio.ScalarImage(tensor=img, affine=MedImage.affine_matrix), 
                                        mask=tio.LabelMap(tensor=y_true, affine=MedImage.affine_matrix)))
            return MedImage.create(aug['img'].data), MedMask.create(aug['mask'].data)

        aug = self.aug(tio.Subject(img=tio.ScalarImage(tensor=img)))
        return MedImage.create(aug['img'].data), y_true


# %% ../nbs/03_vision_augment.ipynb 7
def do_pad_or_crop(o, target_shape, padding_mode, mask_name, dtype=torch.Tensor):
    #TODO:refactorize
    pad_or_crop = tio.CropOrPad(target_shape=target_shape, padding_mode=padding_mode, mask_name=mask_name)
    return dtype(pad_or_crop(o))

# %% ../nbs/03_vision_augment.ipynb 8
class PadOrCrop(DisplayedTransform):
    """Resize image using TorchIO `CropOrPad`."""
    
    order = 0

    def __init__(self, size, padding_mode=0, mask_name=None):
        if not is_listy(size): 
            size = [size, size, size]
        self.pad_or_crop = tio.CropOrPad(target_shape=size,
                                    padding_mode=padding_mode, 
                                    mask_name=mask_name)

    def encodes(self, o: (MedImage, MedMask)):
        return type(o)(self.pad_or_crop(o))

# %% ../nbs/03_vision_augment.ipynb 9
class ZNormalization(DisplayedTransform):
    """Apply TorchIO `ZNormalization`."""

    order = 0

    def __init__(self, masking_method=None, channel_wise=True):
        self.z_normalization = tio.ZNormalization(masking_method=masking_method)
        self.channel_wise = channel_wise

    def encodes(self, o: MedImage):
        if self.channel_wise:
            o = torch.stack([self.z_normalization(c[None])[0] for c in o])
        else: o = self.z_normalization(o) 

        return MedImage.create(o)

    def encodes(self, o: MedMask):
        return o

# %% ../nbs/03_vision_augment.ipynb 10
class BraTSMaskConverter(DisplayedTransform):
    '''Convert BraTS masks.'''

    order=1

    def encodes(self, o:(MedImage)): return o

    def encodes(self, o:(MedMask)):
        o = torch.where(o==4, 3., o)
        return MedMask.create(o)

# %% ../nbs/03_vision_augment.ipynb 11
class BinaryConverter(DisplayedTransform):
    '''Convert to binary mask.'''

    order=1

    def encodes(self, o: MedImage): 
        return o

    def encodes(self, o: MedMask):
        o = torch.where(o>0, 1., 0)
        return MedMask.create(o)

# %% ../nbs/03_vision_augment.ipynb 12
class RandomGhosting(DisplayedTransform):
    """Apply TorchIO `RandomGhosting`."""
    
    split_idx, order = 0, 1

    def __init__(self, intensity=(0.5, 1), p=0.5):
        self.add_ghosts = tio.RandomGhosting(intensity=intensity, p=p)

    def encodes(self, o: MedImage):
        return MedImage.create(self.add_ghosts(o))

    def encodes(self, o: MedMask):
        return o

# %% ../nbs/03_vision_augment.ipynb 13
class RandomSpike(DisplayedTransform):
    '''Apply TorchIO `RandomSpike`.'''
    
    split_idx,order=0,1

    def __init__(self, num_spikes=1, intensity=(1, 3), p=0.5):
        self.add_spikes = tio.RandomSpike(num_spikes=num_spikes, intensity=intensity, p=p)

    def encodes(self, o:MedImage): 
        return MedImage.create(self.add_spikes(o))
        
    def encodes(self, o:MedMask):
        return o

# %% ../nbs/03_vision_augment.ipynb 14
class RandomNoise(DisplayedTransform):
    '''Apply TorchIO `RandomNoise`.'''

    split_idx,order=0,1

    def __init__(self, mean=0, std=(0, 0.25), p=0.5):
        self.add_noise = tio.RandomNoise(mean=mean, std=std, p=p)

    def encodes(self, o: MedImage): 
        return MedImage.create(self.add_noise(o))
    
    def encodes(self, o: MedMask):
        return o

# %% ../nbs/03_vision_augment.ipynb 15
class RandomBiasField(DisplayedTransform):
    '''Apply TorchIO `RandomBiasField`.'''

    split_idx,order=0,1

    def __init__(self, coefficients=0.5, order=3, p=0.5):
        self.add_biasfield = tio.RandomBiasField(coefficients=coefficients, order=order, p=p)

    def encodes(self, o: MedImage): 
        return MedImage.create(self.add_biasfield(o))
        
    def encodes(self, o: MedMask):
        return o

# %% ../nbs/03_vision_augment.ipynb 16
class RandomBlur(DisplayedTransform):
    '''Apply TorchIO `RandomBiasField`.'''

    split_idx,order=0,1

    def __init__(self, std=(0, 2), p=0.5):
        self.add_blur = tio.RandomBlur(std=std, p=p)
        
    def encodes(self, o: MedImage): 
        return MedImage.create(self.add_blur(o))
    
    def encodes(self, o: MedMask):
        return o

# %% ../nbs/03_vision_augment.ipynb 17
class RandomGamma(DisplayedTransform):
    '''Apply TorchIO `RandomGamma`.'''


    split_idx,order=0,1

    def __init__(self, log_gamma=(-0.3, 0.3), p=0.5):
        self.add_gamma = tio.RandomGamma(log_gamma=log_gamma, p=p)

    def encodes(self, o: MedImage): 
        return MedImage.create(self.add_gamma(o))
    
    def encodes(self, o: MedMask):
        return o

# %% ../nbs/03_vision_augment.ipynb 18
class RandomMotion(DisplayedTransform):
    """Apply TorchIO `RandomMotion`."""

    split_idx, order = 0, 1

    def __init__(
            self, 
            degrees=10, 
            translation=10, 
            num_transforms=2, 
            image_interpolation='linear', 
            p=0.5
        ):
        self.add_motion = tio.RandomMotion(
            degrees=degrees, 
            translation=translation, 
            num_transforms=num_transforms, 
            image_interpolation=image_interpolation, 
            p=p
        )

    def encodes(self, o: MedImage):
        return MedImage.create(self.add_motion(o))

    def encodes(self, o: MedMask):
        return o

# %% ../nbs/03_vision_augment.ipynb 20
class RandomElasticDeformation(CustomDictTransform):
    """Apply TorchIO `RandomElasticDeformation`."""

    def __init__(self, num_control_points=7, max_displacement=7.5,
                 image_interpolation='linear', p=0.5):
                     
        super().__init__(tio.RandomElasticDeformation(
            num_control_points=num_control_points,
            max_displacement=max_displacement,
            image_interpolation=image_interpolation,
            p=p))

# %% ../nbs/03_vision_augment.ipynb 21
class RandomAffine(CustomDictTransform):
    """Apply TorchIO `RandomAffine`."""

    def __init__(self, scales=0, degrees=10, translation=0, isotropic=False,
                 image_interpolation='linear', default_pad_value=0., p=0.5):
                     
        super().__init__(tio.RandomAffine(
            scales=scales,
            degrees=degrees,
            translation=translation,
            isotropic=isotropic,
            image_interpolation=image_interpolation,
            default_pad_value=default_pad_value,
            p=p))

# %% ../nbs/03_vision_augment.ipynb 22
class RandomFlip(CustomDictTransform):
    """Apply TorchIO `RandomFlip`."""

    def __init__(self, axes='LR', p=0.5):
        super().__init__(tio.RandomFlip(axes=axes, flip_probability=p))

# %% ../nbs/03_vision_augment.ipynb 23
class OneOf(CustomDictTransform):
    """Apply only one of the given transforms using TorchIO `OneOf`."""

    def __init__(self, transform_dict, p=1):
        super().__init__(tio.OneOf(transform_dict, p=p))
