custom_transforms.py

import math
from typing import List, Tuple, Optional, Dict
import os

import torch
from torch import Tensor
import numpy as np
import PIL.Image
import random
from io import BytesIO
import cv2
import numpy as np

from torchvision.transforms import functional as F, InterpolationMode
import torchvision.transforms as T

__all__ = ["AutoAugmentPolicy", "AutoAugment", "RandAugment", "TrivialAugmentWide", "AugMix"]

def get_dimensions(img):
    height, width = F.get_image_size(img)
    channels = F.get_image_num_channels(img)
    return channels, height, width

def cutout(img, pad_size, replace=0):
    """Apply cutout (https://arxiv.org/abs/1708.04552) to image. 
    
    ### (PyTorch implementation of Google's big_vision cutout) ###
    
    This operation applies a (2*pad_size x 2*pad_size) mask of zeros to
    a random location within `img`. The pixel values filled in will be of the
    value `replace`. The located where the mask will be applied is randomly
    chosen uniformly over the whole image.
    Args:
        image: A PIL image
        pad_size: Specifies how big the zero mask that will be generated is that
        is applied to the image. The mask will be of size
        (2*pad_size x 2*pad_size).
        replace: What pixel value to fill in the image in the area that has
        the cutout mask applied to it.
    Returns:
        A PIL image of type uint8.
    """
    convert_back=False
    if F._is_pil_image(img):
        img = F.pil_to_tensor(img) # convert to tensor for pytorch operations
        convert_back=True
    assert img.dtype == torch.uint8, "PIL to tensor image is expected to have torch.unit8 as dtype."
    channels, height, width = get_dimensions(img)
    cutout_center_height = torch.randint(low=0, high=height, size=(1,)).item()
    cutout_center_width = torch.randint(low=0, high=width, size=(1,)).item()

    lower_pad = max(0, cutout_center_height - pad_size)
    upper_pad = max(0, height - cutout_center_height - pad_size)
    left_pad = max(0, cutout_center_width - pad_size)
    right_pad = max(0, width - cutout_center_width - pad_size)

    cutout_shape = (height - (lower_pad + upper_pad),
                    width - (left_pad + right_pad)) # cutout this shape
    padding_dims = (left_pad, right_pad, upper_pad, lower_pad)
    cutout_mask = torch.nn.functional.pad(
        torch.zeros(cutout_shape, dtype=img.dtype, device=img.device),
        padding_dims, value=1
    )
    cutout_mask = cutout_mask.unsqueeze(dim=0)
    cutout_mask = torch.tile(cutout_mask, (channels,1,1))
    #replacement = torch.ones_like(img, dtype=torch.float32) * replace[0]
    #replacement = replacement.to(torch.uint8)
    img = torch.where(
        cutout_mask==0, # condition.
        torch.ones_like(img, dtype=img.dtype, device=img.device) * replace, # If true
        #replacement,
        img # If condition is false
    )
    if convert_back:
        return F.to_pil_image(img)
    else:
        return img

def solarize_add(img, addition=0, threshold=128):
    """
    For each pixel in the image less than threshold
    we add 'addition' amount to it and then clip the
    pixel value to be between 0 and 255. The value
    of 'addition' is between -128 and 128.
    
    ### Re-implementation of Google's big_vision in PyTorch ###
    """
    convert_back=False
    if F._is_pil_image(img):
        img = F.pil_to_tensor(img) # convert to tensor for pytorch operations
        convert_back=True
    assert img.dtype == torch.uint8, "PIL to tensor image is expected to have torch.unit8 as dtype."
    added_img = img.to(torch.int) + addition
    added_img = torch.clamp(added_img, min=0,max=255)
    added_img = added_img.to(img.dtype)
    img = torch.where(
        img < threshold, # condition
        added_img, # if true
        img # if false
    )
    if convert_back:
        return F.to_pil_image(img)
    else:
        return img

def chroma_drop(img):
    img = img.convert("YCbCr")
    Y, Cb, Cr = img.split()
    if torch.rand(1).item() > 0.5:
        Cr = Cr.point(lambda i: 128)
    else:
        Cb = Cb.point(lambda i: 128)
    img = PIL.Image.merge("YCbCr", (Y, Cb, Cr))
    return img.convert("RGB")

def auto_saturation_separate(img):
    img = img.convert("YCbCr")
    Y, Cb, Cr = img.split()
    Cbmin, Cbmax = Cb.getextrema()
    Crmin, Crmax = Cr.getextrema()
    Cmin = min(Cbmin, Crmin)
    Cmax = max(Cbmax, Crmax)
    Cb = Cb.point(lambda i: ((i-128) / (Cmax - 128) * 127 + 128 if Cmax > 128 else i) if i>127 \
        else ((i - Cmin) / (127 - Cmin) * 127) if Cmin<127 else i) # scale >127 and else separately (they represent different hue)
    #Cb = Cb.point(lambda i: (i-Cbmin) / (Cbmax - Cbmin) * 255)
    Cr = Cr.point(lambda i: ((i-128) / (Cmax - 128) * 127 + 128 if Cmax > 128 else i) if i>127 \
        else ((i - Cmin) / (127 - Cmin) * 127) if Cmin<127 else i)
    #Cr = Cr.point(lambda i: (i-Crmin) / (Crmax - Crmin) * 255)
    img = PIL.Image.merge("YCbCr", (Y, Cb, Cr))
    return img.convert("RGB")


def auto_saturation(img):
    img = img.convert("YCbCr")
    Y, Cb, Cr = img.split()
    Cbmin, Cbmax = Cb.getextrema()
    Crmin, Crmax = Cr.getextrema()
    Cmin = min(Cbmin, Crmin)
    Cmax = max(Cbmax, Crmax)
    Cb = Cb.point(lambda i: (i-Cmin) / (Cmax - Cmin) * 255 if (Cmax - Cmin) != 0 else i)
    Cr = Cr.point(lambda i: (i-Cmin) / (Cmax - Cmin) * 255 if (Cmax - Cmin) != 0 else i)
    img = PIL.Image.merge("YCbCr", (Y, Cb, Cr))
    return img.convert("RGB")

def _apply_op(
    img: Tensor, op_name: str, magnitude: float, interpolation: InterpolationMode, fill: Optional[List[float]]
):
    if op_name == "ShearX":
        # magnitude should be arctan(magnitude)
        # official autoaug: (1, level, 0, 0, 1, 0)
        # https://github.com/tensorflow/models/blob/dd02069717128186b88afa8d857ce57d17957f03/research/autoaugment/augmentation_transforms.py#L290
        # compared to
        # torchvision:      (1, tan(level), 0, 0, 1, 0)
        # https://github.com/pytorch/vision/blob/0c2373d0bba3499e95776e7936e207d8a1676e65/torchvision/transforms/functional.py#L976
        img = F.affine(
            img,
            angle=0.0,
            translate=[0, 0],
            scale=1.0,
            shear=[math.degrees(math.atan(magnitude)), 0.0],
            interpolation=interpolation,
            fill=fill,
            center=[0, 0],
        )
    elif op_name == "ShearY":
        # magnitude should be arctan(magnitude)
        # See above
        img = F.affine(
            img,
            angle=0.0,
            translate=[0, 0],
            scale=1.0,
            shear=[0.0, math.degrees(math.atan(magnitude))],
            interpolation=interpolation,
            fill=fill,
            center=[0, 0],
        )
    elif op_name == "TranslateX":
        img = F.affine(
            img,
            angle=0.0,
            translate=[int(magnitude), 0],
            scale=1.0,
            interpolation=interpolation,
            shear=[0.0, 0.0],
            fill=fill,
        )
    elif op_name == "TranslateY":
        img = F.affine(
            img,
            angle=0.0,
            translate=[0, int(magnitude)],
            scale=1.0,
            interpolation=interpolation,
            shear=[0.0, 0.0],
            fill=fill,
        )
    elif op_name == "Rotate":
        img = F.rotate(img, magnitude, interpolation=interpolation, fill=fill)
    elif op_name == "Brightness":
        img = F.adjust_brightness(img, 1.0 + magnitude)
    elif op_name == "Color":
        img = F.adjust_saturation(img, 1.0 + magnitude)
    elif op_name == "Contrast":
        img = F.adjust_contrast(img, 1.0 + magnitude)
    elif op_name == "Sharpness":
        img = F.adjust_sharpness(img, 1.0 + magnitude)
    elif op_name == "Posterize":
        img = F.posterize(img, int(magnitude))
    elif op_name == "Solarize":
        img = F.solarize(img, magnitude)
    elif op_name == "AutoContrast":
        img = F.autocontrast(img)
    elif op_name == "Equalize":
        img = F.equalize(img)
    elif op_name == "Invert":
        img = F.invert(img)
    elif op_name == "Identity":
        pass
    elif op_name == 'Cutout': # added
        img = cutout(img, int(magnitude), replace=fill)
    elif op_name == "SolarizeAdd": # added
        img = solarize_add(img, int(magnitude))
    elif op_name == "Grayscale": # added v2
        img = F.to_grayscale(img, num_output_channels=3)
    elif op_name == "ChromaDrop": #
        img = chroma_drop(img)
    elif op_name == "AutoSaturation":
        #img = auto_saturation(img)
        img = auto_saturation(img) # dct-equivalent
    elif op_name == "AutoSaturation_old": # for compatibility purposes
        img = auto_saturation(img)
    elif op_name == "Rotate90": # magnitude is +- 90
        img = F.rotate(img, magnitude, interpolation=interpolation, fill=fill)
    else:
        raise ValueError(f"The provided operator {op_name} is not recognized.")
    return img



class RandAugment_bv(torch.nn.Module):
    r"""RandAugment data augmentation method based on
    `"RandAugment: Practical automated data augmentation with a reduced search space"
    <https://arxiv.org/abs/1909.13719>`_.

    ### Re-implementation of Google's Big Vision randaugment in PyTorch ###

    If the image is torch Tensor, it should be of type torch.uint8, and it is expected
    to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
    If img is PIL Image, it is expected to be in mode "L" or "RGB".

    Args:
        num_ops (int): Number of augmentation transformations to apply sequentially.
        magnitude (int): Magnitude for all the transformations.
        num_magnitude_bins (int): The number of different magnitude values.
        interpolation (InterpolationMode): Desired interpolation enum defined by
            :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``.
            If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported.
        fill (sequence or number, optional): Pixel fill value for the area outside the transformed
            image. If given a number, the value is used for all bands respectively.
    """

    def __init__(
        self,
        num_ops: int = 2,
        magnitude: int = 10,
        num_magnitude_bins: int = 11,
        interpolation: InterpolationMode = InterpolationMode.NEAREST,
        fill: Optional[List[float]] = None,
        ops_list = ["AutoContrast", "Equalize", "Invert", "Rotate", "Posterize", "Solarize", "SolarizeAdd", "Color", "Contrast", "Brightness",
                        "Sharpness", "ShearX", "ShearY", "Cutout", "TranslateX", "TranslateY"]
    ) -> None:
        super().__init__()
        self.num_ops = num_ops
        self.magnitude = magnitude
        self.num_magnitude_bins = num_magnitude_bins
        self.interpolation = interpolation
        self.fill = fill
        if ops_list==None:
            self.ops_list = ["AutoContrast", "Equalize", "Invert", "Rotate", "Posterize", "Solarize", "SolarizeAdd", "Color", "Contrast", "Brightness",
                        "Sharpness", "ShearX", "ShearY", "Cutout", "TranslateX", "TranslateY"]
        else:
            self.ops_list = ops_list

    def _augmentation_space(self, num_bins: int, image_size: Tuple[int, int]) -> Dict[str, Tuple[Tensor, bool]]:
        return {
            # op_name: (magnitudes, signed)
            #"Identity": (torch.tensor(0.0), False), not needed
            "AutoContrast": (torch.tensor(0.0), False),
            "Equalize": (torch.tensor(0.0), False),
            "Invert": (torch.tensor(0.0), False), # added
            "Rotate": (torch.linspace(0.0, 30.0, num_bins), True),
            "Posterize": (8 - (torch.arange(num_bins) / ((num_bins - 1) / 4)).round().int(), False),
            "Solarize": (torch.linspace(255.0, 0.0, num_bins), False),
            "SolarizeAdd": (torch.linspace(0, 110, num_bins), False), # added
            "Color": (torch.linspace(0.0, 0.9, num_bins), True),
            "Contrast": (torch.linspace(0.0, 0.9, num_bins), True),
            "Brightness": (torch.linspace(0.0, 0.9, num_bins), True),
            "Sharpness": (torch.linspace(0.0, 0.9, num_bins), True),
            "ShearX": (torch.linspace(0.0, 0.3, num_bins), True),
            "ShearY": (torch.linspace(0.0, 0.3, num_bins), True),
            "Cutout": (torch.linspace(0, 40, num_bins), False), #added
            "TranslateX": (torch.linspace(0.0, 150.0 / 336.0 * image_size[1], num_bins), True),
            "TranslateY": (torch.linspace(0.0, 150.0 / 336.0 * image_size[0], num_bins), True),
            "Grayscale": (torch.tensor(0.0), False),
            "ChromaDrop": (torch.tensor(0.0), False),
            "AutoSaturation": (torch.tensor(0.0), False),
            "AutoSaturation_old": (torch.tensor(0.0), False),
            "Rotate90": (torch.tensor(90.0), True),
        }


    def forward(self, img: Tensor) -> Tensor:
        """
            img (PIL Image or Tensor): Image to be transformed.

        Returns:
            PIL Image or Tensor: Transformed image.
        """
        fill = self.fill
        channels, height, width = get_dimensions(img)
        #if isinstance(img, Tensor):
        #    if isinstance(fill, (int, float)):
        #        fill = [float(fill)] * channels
        #    elif fill is not None:
        #        fill = [float(f) for f in fill]

        op_meta = self._augmentation_space(self.num_magnitude_bins, (height, width))
        for _ in range(self.num_ops):
            op_index = int(torch.randint(len(self.ops_list), (1,)).item())
            op_name = list(self.ops_list)[op_index]
            magnitudes, signed = op_meta[op_name]
            magnitude = float(magnitudes[self.magnitude].item()) if magnitudes.ndim > 0 else 0.0
            if signed and torch.randint(2, (1,)):
                magnitude *= -1.0
            img = _apply_op(img, op_name, magnitude, interpolation=self.interpolation, fill=fill)

        return img


    def __repr__(self) -> str:
        s = (
            f"{self.__class__.__name__}("
            f"num_ops={self.num_ops}"
            f", magnitude={self.magnitude}"
            f", num_magnitude_bins={self.num_magnitude_bins}"
            f", interpolation={self.interpolation}"
            f", fill={self.fill}"
            f")"
        )
        return s


class ToTensor_range(torch.nn.Module):
    r"""
    Converts PIL image to Tensor into a specified range

    Args:
        val_min = minimum value after convert
        val_max = maximum value after convert
        dtype = dtype after convert (default=torch.float32)

    Returns:
        Converted Torch Tensor
    """

    def __init__(
        self,
        val_min: float = -1.,
        val_max: float = 1.,
        dtype = torch.float32,
    ) -> Tensor:
        super().__init__()
        self.val_min = val_min
        self.val_max = val_max
        self.dtype = dtype

    def forward(self, img) -> Tensor:
        """
            img (PIL Image): Image to be transformed.

        Returns:
            Tensor: Converted Image
        """
        #assert F._is_pil_image(img), "Input should be a PIL image (ToTensor_range transform)"
        if F._is_pil_image(img):
            img = F.to_tensor(img) # to_tensor normalizes data to (0,1)
        img = img.to(self.dtype) # convert dtype
        img = self.val_min + (img * (self.val_max - self.val_min)) # scale to val_min to val_max

        return img

    def __repr__(self) -> str:
        s = (
            f"{self.__class__.__name__}("
            f"val_min={self.val_min}"
            f", val_max={self.val_max}"
            f", dtype={self.dtype}"
            f")"
        )
        return s

def apply_PILJPEG(img, quality):
    buffer = BytesIO()
    img.save(buffer, format="JPEG", quality=quality)
    buffer.seek(0) # move pointer to 0 so we can read them
    img = PIL.Image.open(buffer).convert("RGB")
    return img

def apply_cv2JPEG(img, quality):
    # convert PIL image to cv2 image
    img_cv2 = np.array(img)
    img_cv2 = img_cv2[:,:,::-1]
    encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), quality]
    result, encimg = cv2.imencode('.jpg', img_cv2, encode_param)
    decimg = cv2.imdecode(encimg, 1)
    return PIL.Image.fromarray(decimg[:,:,::-1])

def apply_randomJPEG(img, quality):
    if random.random() < 0.5:
        img = apply_PILJPEG(img, quality) # randomly apply PIL or CV2
    else:
        img = apply_cv2JPEG(img, quality)
    return img

def resize_with_random_intpl(img, size):
    """
    Perform resizing with random interpolation
    """
    #intp_list = [InterpolationMode.NEAREST, InterpolationMode.BILINEAR, InterpolationMode.BICUBIC, InterpolationMode.LANCZOS, InterpolationMode.HAMMING, InterpolationMode.BOX]
    intp_list = [InterpolationMode.BILINEAR, InterpolationMode.BICUBIC]
    #interp_idx = random.randint(0, len(intp_list)-1)
    interp = random.choice(intp_list)
    # random interpolation somehow doesn't work
    img = F.resize(img, size, interpolation=interp)
    return img

class RandomResizeWithRandomIntpl(torch.nn.Module):
    r"""
    Reads PIL Image. Resizes with random interpolation. Returns torch tensor.
    """

    def __init__(
        self,
        size_range: int=(112,448),
    ) -> Tensor:
        super().__init__()
        self.size_range = size_range

    def forward(self, img) -> Tensor:
        """
        Args:
            img: PIL image to be transformed.

        Returns:
            Tensor: Converted Image
        """
        assert F._is_pil_image(img), "Input should be a PIL image (RandomResizeWithRandomIntpl transform)"
        # add resize
        img = resize_with_random_intpl(img, random.randint(self.size_range[0], self.size_range[1]))
        return img

    def __repr__(self) -> str:
        s = (
            f"{self.__class__.__name__}()"
            f" size_range={self.size_range}"
            f")"
        )

class ResizeWithRandomIntpl(torch.nn.Module):
    r"""
    Reads PIL Image. Resizes with random interpolation. Returns torch tensor.
    """

    def __init__(
        self,
        size: int,
    ) -> Tensor:
        super().__init__()
        self.size = size

    def forward(self, img) -> Tensor:
        """
        Args:
            img: PIL image to be transformed.

        Returns:
            Tensor: Converted Image
        """
        assert F._is_pil_image(img), "Input should be a PIL image (ResizeWithRandomIntpl transform)"
        # add resize
        img = resize_with_random_intpl(img, self.size)
        return img

    def __repr__(self) -> str:
        s = (
            f"{self.__class__.__name__}("
            f" size={self.size}"
            f")"
        )
        return s

class RRCWithRandomIntpl(T.RandomResizedCrop):
    r"""
    Reads PIL Image. Randomly resized crop with random interpolation. Returns torch tensor.
    """

    def __init__(
        self,
        size: int,
        scale: Tuple[float, float] = (0.08, 1.0),
        ratio: Tuple[float, float] = (3./4., 4./3.),
    ) -> Tensor:
        super().__init__(size=size, scale=scale, ratio=ratio)
        self.size = size
        self.scale = scale
        self.ratio = ratio
        self.intp_list=[InterpolationMode.NEAREST, InterpolationMode.BILINEAR, InterpolationMode.BICUBIC, InterpolationMode.LANCZOS, InterpolationMode.HAMMING, InterpolationMode.BOX]

    def forward(self, img) -> Tensor:
        """
        Args:
            img: PIL image to be transformed.

        Returns:
            Tensor: Converted Image
        """
        assert F._is_pil_image(img), "Input should be a PIL image (RRCWithRandomIntpl transform)"
        # add resize
        i, j, h, w = self.get_params(img, self.scale, self.ratio)
        #interp_idx = random.randint(0, len(self.intp_list)-1)
        interp = random.choice(self.intp_list) # somehow doesn't work. Gives me error: TypeError: resized_crop() got multiple values for argument 'interpolation'
        return F.resized_crop(img, i, j, h, w, self.size, interpolation=interp)

    def __repr__(self) -> str:
        s = (
            f"{self.__class__.__name__}("
            f" size={self.size}"
            f", scale={self.scale}"
            f", ratio={self.ratio}"
            f")"
        )
        return s

class JPEGinMemory(torch.nn.Module):
    r"""
    Reads PIL Image. Compress JPEG in memory. Returns PIL Image.

    """

    def __init__(
        self,
        quality_range = (30, 100),
        method: str = "cv,pil",
        dtype = torch.float32,
    ) -> Tensor:
        super().__init__()
        self.quality_range = quality_range
        self.method = method.lower().split(',')
        self.dtype = dtype

    def forward(self, img) -> Tensor:
        """
        Args:
            img: PIL image to be transformed.jdt

        Returns:
            Tensor: Converted Image
        """
        assert F._is_pil_image(img), "Input should be a PIL image (ResizeAndJPEGinMemory transform)"
        if "cv" in self.method and "pil" in self.method:
            img = apply_randomJPEG(img, random.randint(self.quality_range[0], self.quality_range[1]))
        elif "cv" in self.method:
            img = apply_cv2JPEG(img, random.randint(self.quality_range[0], self.quality_range[1]))
        elif "pil" in self.method:
            img = apply_PILJPEG(img, random.randint(self.quality_range[0], self.quality_range[1]))
        return img

    def __repr__(self) -> str:
        s = (
            f"{self.__class__.__name__}("
            f", quality_range={self.quality_range}"
            f", dtype={self.dtype}"
            f")"
        )
        return s

class ResizeAndJPEGinMemory(torch.nn.Module):
    r"""
    Reads PIL Image. Resizes and compresses to JPEG in memory. Returns torch tensor.

    """

    def __init__(
        self,
        size: int,
        quality: int = 95,
        method: str = "cv,pil",
        dtype = torch.float32,
    ) -> Tensor:
        super().__init__()
        self.size = size
        self.quality = quality
        self.method = method.lower().split(',')
        self.dtype = dtype

    def forward(self, img) -> Tensor:
        """
        Args:
            img: PIL image to be transformed.

        Returns:
            Tensor: Converted Image
        """
        assert F._is_pil_image(img), "Input should be a PIL image (ResizeAndJPEGinMemory transform)"
        # add resize
        img = F.resize(img, self.size, interpolation=InterpolationMode.BILINEAR) # this is the right way to resize! If torchvision updates, make sure that this resizes the smaller side to the specified size and keeps the aspect ratio
        if "cv" in self.method and "pil" in self.method:
            img = apply_randomJPEG(img, self.quality)
        elif "cv" in self.method:
            img = apply_cv2JPEG(img, self.quality)
        elif "pil" in self.method:
            img = apply_PILJPEG(img, self.quality)
        return img

    def __repr__(self) -> str:
        s = (
            f"{self.__class__.__name__}("
            f" size={self.size}"
            f", quality={self.quality}"
            f", dtype={self.dtype}"
            f")"
        )
        return s
    
class StochasticJPEG(torch.nn.Module):
    r"""
    Stochastically applies multiple JPEG compression and resizing to an image.
    """

    def __init__(
        self,
        size: int, # final output size
        quality: Tuple[int, int] = (50, 100), # quality range
        num_jpeg: Tuple[int, int] = (1, 5), # number of jpegs to apply
        jpeg_p: float = 0.5, # probability of applying JPEG compression
        rrc_p: float = 0.5, # probability of applying random resized crop
        rrc_scale: Tuple[float, float] = (0.75, 1.0), # random resize crop scale
        rrc_ratio: Tuple[float, float] = (3./4., 4./3.), # random resize crop ratio
        no_rrc: bool = False, # if True, no random resized crop is applied
        dtype: type = torch.float32,
    ) -> Tensor:
        """
        Initialize the CustomTransforms class.

        Args:
            size (int): The final output size.
            quality (Tuple[int, int]): The quality range as a tuple of two integers.
            num_jpeg (Tuple[int, int]): The number of jpegs to apply as a tuple of two integers.
            p (float): The probability of applying the transform.
            rrc_scale (Tuple[float, float]): The random resize crop scale as a tuple of two floats.
            rrc_ratio (Tuple[float, float]): The random resize crop ratio as a tuple of two floats.
            no_rrc (bool): If True, no random resized crop is applied.
            dtype (type): The data type of the tensor.

        Returns:
            Tensor: The initialized CustomTransforms object.
        """
        super().__init__()
        self.size = size
        self.quality = quality
        self.num_jpeg = num_jpeg
        self.jpeg_p = jpeg_p
        self.rrc_p = rrc_p
        self.rrc = torch.nn.Identity() if no_rrc else T.RandomResizedCrop(size=size, scale=rrc_scale, ratio=rrc_ratio, interpolation=InterpolationMode.BILINEAR)
        self.dtype = dtype

    def forward(self, img) -> Tensor:
        """
        Args:
            img: PIL image to be transformed.

        Returns:
            Tensor: Converted Image
        """
        assert F._is_pil_image(img), "Input should be a PIL image (StochasticJPEG transform)"

        # randomly sample p
        count = self.num_jpeg[0]
        for _ in range(self.num_jpeg[0]): # apply min number of jpegs and RRC first
            img = self.rrc(img)
            img = apply_randomJPEG(img, random.randint(self.quality[0], self.quality[1]))
        
        while count < self.num_jpeg[1]:
            if random.random() < self.p: # apply more jpegs with set probability.
                img = self.rrc(img)
                img = apply_randomJPEG(img, random.randint(self.quality[0], self.quality[1]))
                count += 1
            else:
                break
        
        return img

class RandomJPEG(torch.nn.Module):
    """
    Randomly applies JPEG
    Args:
        quality: tuple of quality value range for JPEG
        p: probability of applying JPEG
    """
    def __init__(
        self,
        quality_list: tuple = (30, 100),
        p: float = 0.5,
    ):
        super().__init__()
        self.quality_list = quality_list
        self.p = p
    
    def forward(self, img):
        if random.random() < self.p:
            img = apply_randomJPEG(img, random.randint(self.quality_list[0], self.quality_list[1]))
        return img

class RandomGaussianBlur(torch.nn.Module):
    """
    Randomly applies Gaussian Blur
    Args:
        p: probability of applying JPEG
        sigma: tuple of sigma values for Gaussian Blur
    """
    def __init__(
        self,
        p: float = 0.5,
        sigma: Tuple[float, float] = (0.0, 3.0),
    ):
        super().__init__()
        self.p = p
        self.sigma = sigma
    
    def forward(self, img):
        if random.random() < self.p:
            sigma=random.uniform(self.sigma[0], self.sigma[1])
            kernel_size=1+2*round(sigma*4.0) # default sigma used in scipy (https://github.com/scipy/scipy/blob/v1.13.1/scipy/ndimage/_filters.py#L286-L390)
            img = F.gaussian_blur(img, kernel_size=kernel_size, sigma=sigma)
        return img

class RandomPaddingAndResize(torch.nn.Module):
    r"""
    Reads PIL Image. Randomly applies padding, and resize it back to original resolution.

    """

    def __init__(
        self,
        pad_percentage_range = (0.1, 0.1), # random padding percentage for x (width) and y (height)
        padding_value_range = (0, 255), # random padding value range
    ) -> Tensor:
        super().__init__()
        self.pad_percentage_range = pad_percentage_range
        self.padding_value_range = padding_value_range

    def forward(self, img) -> Tensor:
        """
        Args:
            img: PIL image to be transformed.jdt

        Returns:
            Tensor: Converted Image
        """
        assert F._is_pil_image(img), "Input should be a PIL image (ResizeAndJPEGinMemory transform)"
        original_size = img.size
        pad_x_l = random.uniform(0, self.pad_percentage_range[0]/2) # x-axis random padding ratio (left)
        pad_x_r = random.uniform(0, self.pad_percentage_range[0]/2) # x-axis random padding ratio (right)
        pad_y_l = random.uniform(0, self.pad_percentage_range[1]/2) # y-axis random padding ratio (left)
        pad_y_r = random.uniform(0, self.pad_percentage_range[1]/2) # y-axis random padding ratio (right)
        pad_fill = random.randint(int(self.padding_value_range[0]), int(self.padding_value_range[1])) # random padding fill value
        img = F.pad(img, (int(pad_x_l*img.size[0]), int(pad_y_l*img.size[1]), int(pad_x_r*img.size[0]), int(pad_y_r*img.size[1])), fill=pad_fill, padding_mode='constant')
        img = F.resize(img, original_size, interpolation=InterpolationMode.BILINEAR)
        return img

    def __repr__(self) -> str:
        s = (
            f"{self.__class__.__name__}("
            f", pad_percentage_range={self.pad_percentage_range}"
            f", padding_value_range={self.padding_value_range}"
            f")"
        )
        return s

class RandomCutout(T.RandomErasing):
    r"""
    Random cutout with random numbers
    """
    def __init__(
        self,
        p=0.5,
        scale=(0.02, 0.33),
        ratio=(0.3, 3.3),
        value_range=(0, 255),
    ):
        super().__init__(p=p, scale=scale, ratio=ratio)
        self.value_range = value_range
    
    def forward(self, img):
        convert_to_pil=False
        if F._is_pil_image(img):
            img = F.pil_to_tensor(img)
            convert_to_pil=True
        if torch.rand(1) < self.p:
            rand_value = random.randint(self.value_range[0], self.value_range[1])
            # cast self.value to script acceptable type
            if isinstance(rand_value, (int, float)):
                rand_value = [float(rand_value)]
            elif isinstance(rand_value, str):
                rand_value = None
            elif isinstance(rand_value, (list, tuple)):
                rand_value = [float(v) for v in rand_value]
            else:
                rand_value = rand_value

            if rand_value is not None and not (len(rand_value) in (1, img.shape[-3])):
                raise ValueError(
                    "If value is a sequence, it should have either a single value or "
                    f"{img.shape[-3]} (number of input channels)"
                )
            x, y, h, w, v = self.get_params(img, self.scale, self.ratio, rand_value)
            img = F.erase(img, x, y, h, w, v)
        if convert_to_pil:
            img = F.to_pil_image(img)
        return img

class RandomVisualization(torch.nn.Module):
    r"""
    Randomly visualizes the fully augmented images by saving them at a specified directory.
    """
    def __init__(
        self,
        save_dir: str = "/nfs/turbo/coe-ahowens-nobackup/jespark/visualizations/fake_img",
        save_p: float = 0.01,
        max_imgs: int = 500,
        overwrite: bool = False,
    ) -> None:
        super().__init__()
        self.save_dir = save_dir
        self.save_p = save_p
        self.max_imgs = max_imgs
        self.overwrite = overwrite
        self.skip_namecheck=False

    def next_available_filename(self, save_dir, max_imgs):
        # Returns next available filename
        # image format = visualization_{03d}_{i}.png, i=[0, max_imgs)
        # let's not make it overwrite
        imgs = os.listdir(save_dir)
        imgs_list = [int(img.split("_")[-1].split(".")[0]) for img in imgs]
        random_int = random.randint(0, 999)
        if len(imgs_list) >= max_imgs:
            if self.overwrite:
                return random.choice(imgs) # overwrite random file from imgs
            else:
                self.skip_namecheck=True
                return False
        elif len(imgs_list) > 0:
            next_int = max(imgs_list) + 1
            return f"visualization_{next_int}_{random_int:03d}.png"
        elif len(imgs_list) == 0:
            return f"visualization_0_{random_int:03d}.png"
        else: # uncaught, unexpected situation.
            raise ValueError("Error in next_available_filename")
        
    def forward(self, img) -> Tensor:
        """
        Args:
            img: PIL image to be transformed.

        Returns:
            Tensor: Converted Image
        """
        if not self.skip_namecheck:
            if random.random() < self.save_p:
                os.makedirs(self.save_dir, exist_ok=True)
                filename = self.next_available_filename(self.save_dir, self.max_imgs)
                if filename:
                    img.save(os.path.join(self.save_dir, filename))
        return img

class RandomStateAugmentation(torch.nn.Module):
    r"""
    Randomly applies augmentations given in the input
    """
    def __init__(
        self,
        resize_size=256,
        crop_size=224,
        auglist="JPEGinMemory,RandomResizeWithRandomIntpl,RandomCrop,RandomHorizontalFlip,RandomVerticalFlip,RRCWithRandomIntpl,RandomRotation,RandomTranslate,RandomShear,RandomPadding",
        min_augs='0',
        max_augs='5',
    ):
        """
        auglist: augmentation lists to apply. Input comma-separated string of augmentations.
        min_augs: minimum number of augmentations to apply. (can be comma-separated string to denote per-augmentation minimum)
        max_augs: maximum number of augmentations to apply. (can be comma-separated string to denote per-augmentation maximum)
        """
        super().__init__()
        self.resize_size=resize_size
        self.crop_size=crop_size

        self.auglist = self.parse_auglist(auglist)
        # convert min_augs and max_augs to appropriate format
        min_augs = self.parse_augnums(min_augs)
        max_augs = self.parse_augnums(max_augs)
        if type(min_augs) == list:
            assert type(max_augs) == list, "max_augs should be list if min_augs is list."
            assert len(min_augs) == len(auglist), "min_augs length should be equal to auglist length."
            assert len(max_augs) == len(auglist), "max_augs length should be equal to auglist length."
        # convert min_augs and max_augs to list if they are not
        self.min_augs = [min_augs] * len(self.auglist) if type(min_augs) != list else min_augs
        self.max_augs = [max_augs] * len(self.auglist) if type(max_augs) != list else max_augs
    
    def parse_augnums(self, augsnum):
        # parse min_augs or max_augs. They are expected to be a string of integers, optinally separated by commas.
        augsnum_list = augsnum.split(",")
        if len(augsnum_list) == 1:
            return int(augsnum_list[0])
        else:
            return [int(aug) for aug in augsnum_list]


    def parse_auglist(self, auglist):
        # parse str-comma-separated auglist to list of augmentations
        # default augmentation thoughts: "JPEGinMemory,RandomResizeWithRandomIntpl,RandomCrop,RandomHorizontalFlip,RandomVerticalFlip,RRCWithRandomIntpl,RandomRotation,RandomTranslate,RandomShear,RandomPadding"
        auglist_list = auglist.split(",")
        parsed_list = torch.nn.ModuleList()
        for aug_name in auglist_list:
            if aug_name=='singleJPEG':
                parsed_list.append(ResizeAndJPEGinMemory(size=self.crop_size, quality=95, dtype=torch.float32))
            if aug_name=='StochasticJPEG':
                parsed_list.append(StochasticJPEG(size=self.crop_size, quality=(75, 100), num_jpeg=(1, 5), jpeg_p=0.5, rrc_p=0.5, rrc_scale=(0.75, 1.0), rrc_ratio=(3./4., 4./3.), no_rrc=False, dtype=torch.float32))
            if aug_name=='JPEGinMemory':
                parsed_list.append(JPEGinMemory(quality_range=(75, 100), dtype=torch.float32))
            if aug_name=='RandomResizeWithRandomIntpl':
                parsed_list.append(RandomResizeWithRandomIntpl(size_range=(self.crop_size+1,round(self.crop_size*1.228)))) # should not be smaller; causes issues with Random Crop.
            if aug_name=='RandomCrop':
                parsed_list.append(T.RandomCrop(self.crop_size))
            if aug_name=='RandomHorizontalFlip':
                parsed_list.append(T.RandomHorizontalFlip())
            if aug_name=='RandomVerticalFlip':
                parsed_list.append(T.RandomVerticalFlip())
            if aug_name=='RRCWithRandomIntpl':
                parsed_list.append(RRCWithRandomIntpl(size=self.crop_size, scale=(0.9, 1.0), ratio=(3./4., 4./3.)))
            if aug_name=='RandomRotation':
                parsed_list.append(T.RandomRotation(15, interpolation=InterpolationMode.BILINEAR))
            if aug_name=='RandomTranslate':
                parsed_list.append(T.RandomAffine(degrees=0, translate=(0.1, 0.1), scale=None, shear=None, interpolation=InterpolationMode.BILINEAR))
            if aug_name=='RandomShear':
                parsed_list.append(T.RandomAffine(degrees=0, translate=None, scale=None, shear=(-15, 15, -15, 15), interpolation=InterpolationMode.BILINEAR))
            if aug_name=='RandomPadding' or aug_name=='RandomPaddingAndResize':
                parsed_list.append(RandomPaddingAndResize(pad_percentage_range=(0.1, 0.1), padding_value_range=(0, 255)))
            if aug_name=='RandomCutout':
                parsed_list.append(RandomCutout(p=0.5, scale=(0.02, 0.06), ratio=(0.3, 3.3), value_range=(0, 255)))
        
        return parsed_list

    def generate_randAug_counts(self):
        # Generates random required counts per augmentation
        per_aug_counts = [0] * len(self.auglist)
        for i in range(len(per_aug_counts)):
            per_aug_counts[i] = random.randint(self.min_augs[i], self.max_augs[i])
        return per_aug_counts

    def convert_aug_counts_to_idxList(self, per_aug_counts):
        # convert per augmentation count to list of indices. For example, [1,3,2] = [0,1,1,1,2,2]
        idxList = []
        for i in range(len(per_aug_counts)):
            idxList += [i] * per_aug_counts[i]
        return idxList

    def check_if_complete(self, count, min_augs):
        # not needed
        if type(min_augs) == list:
            min_augs_list = min_augs
        else:
            min_augs_list = [min_augs] * len(self.auglist)
        for i in range(len(min_augs_list)):
            if count[i] < min_augs_list[i]:
                return False
        return True

    def forward(self, img) -> Tensor:
        """
        Args:
            img: PIL image to be transformed.

        Returns:
            Tensor: Converted Image
        """
        assert F._is_pil_image(img), "Input should be a PIL image (RandomStateAugmentation transform)"
        # randomly applies augmentation. Randomly walks through the list of augmentations and applies them. They should be applied at least "min_augs" number of times.
        #count = [0] * len(self.auglist)

        idxList = self.convert_aug_counts_to_idxList(self.generate_randAug_counts())

        while len(idxList) > 0:
            randomIdx = idxList.pop(random.randint(0, len(idxList)-1)) # randomly pop index from idxList
            img = self.auglist[randomIdx](img)
            #count[randomIdx] += 1 # not needed, idxList contains exact amount of augmentations to apply per idx.

        return img
    
class RandomSignRotation(torch.nn.Module):
    r"""
    Randomly rotates the image by given angle. Randomly changes sign.
    """

    def __init__(
        self,
        angle: int,
        interpolation: InterpolationMode = InterpolationMode.BILINEAR,
    ) -> Tensor:
        super().__init__()
        self.angle = angle
        self.interpolation = interpolation

    def forward(self, img) -> Tensor:
        """
        Args:
            img: PIL image to be transformed.

        Returns:
            Tensor: Converted Image
        """
        if random.random() < 0.5:
            angle = -self.angle
        else:
            angle = self.angle
        img = F.rotate(img, angle, interpolation=self.interpolation)
        return img
    
    def __repr__(self) -> str:
        s = (
            f"{self.__class__.__name__}("
            f" angle={self.angle}"
            f", interpolation={self.interpolation}"
            f")"
        )
        return s
    
class RandomResize(torch.nn.Module):
    r"""
    Randomly resizes the input. Either up or downsample and then return it to the original size. Arguments take percentage of resizing (e.g., 0.3 means it can be downsized or upsampled by 30%)
    """
    def __init__(
        self,
        resize_percentage: float,
        interpolation: InterpolationMode = InterpolationMode.BILINEAR,
    ) -> Tensor:
        super().__init__()
        self.resize_percentage = resize_percentage
        self.interpolation = interpolation
    
    def forward(self, img) -> Tensor:
        """
        Args:
            img: PIL image to be transformed.

        Returns:
            Tensor: Converted Image
        """
        if random.random() < 0.5:
            resize_percentage = 1.0 - self.resize_percentage
        else:
            resize_percentage = 1.0 + self.resize_percentage
        original_size_1, original_size_0 = img.size # width, height
        img = F.resize(img, (int(original_size_0*resize_percentage), int(original_size_1*resize_percentage)), interpolation=self.interpolation) # resized height, width
        img = F.resize(img, (original_size_0, original_size_1), interpolation=self.interpolation)
        return img