transforms.py

import random
import numpy as np
from torchvision.transforms import functional as F
from PIL import Image
import torchvision.transforms as T
import torch
from augment import apply_op_both,rand_augment_both

class Compose(object):
    """
    Composes a sequence of transforms.
    Arguments:
        transforms: A list of transforms.
    """
    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, image, label):
        for t in self.transforms:
            image, label = t(image, label)
        return image, label

    def __repr__(self):
        format_string = self.__class__.__name__ + "("
        for t in self.transforms:
            format_string += "\n"
            format_string += "    {0}".format(t)
        format_string += "\n)"
        return format_string


class ToTensor(object):
    def __call__(self, image, target):
        image = F.to_tensor(image)
        target = torch.as_tensor(np.array(target), dtype=torch.int64)
        return image, target

class RandAugment:
    def __init__(self,N,M,prob=1.0,fill=(128,128,128),ignore_value=255):
        self.N=N
        self.M=M
        self.prob=prob
        self.fill=fill
        self.ignore_value=ignore_value
    def __call__(self, image, target):
        return rand_augment_both(image,target,n_ops=self.N,magnitude=self.M,prob=self.prob,fill=self.fill,ignore_value=self.ignore_value)


class Normalize(object):
    """
    Normalizes image by mean and std.
    """
    def __init__(self, mean, std):
        self.mean = mean
        self.std = std

    def __call__(self, image, label):
        image = F.normalize(image, mean=self.mean, std=self.std)
        return image, label

class RandomResize(object):
    def __init__(self, min_size, max_size=None):
        self.min_size = min_size
        if max_size is None:
            max_size = min_size
        self.max_size = max_size

    def __call__(self, image, target):
        size = random.randint(self.min_size, self.max_size)
        image = F.resize(image, size)
        target = F.resize(target, size, interpolation=F.InterpolationMode.NEAREST)
        return image, target

class ColorJitter:
    def __init__(self,brightness=0.2, contrast=0.2, saturation=(0.5,4), hue=0.2):
        self.jitter=T.ColorJitter(brightness=brightness, contrast=contrast, saturation=saturation, hue=hue)
    def __call__(self, image, target):
        image=self.jitter(image)
        return image,target

class AddNoise:#additive gaussian noise
    def __init__(self,factor):
        self.factor=factor
    def __call__(self, image, target):
        factor = random.uniform(0, self.factor)
        image = np.array(image)
        assert(image.dtype==np.uint8)
        gauss = (np.array(torch.randn(*image.shape)) * factor).astype("uint8")
        noisy = (image + gauss).clip(0, 255)
        image = Image.fromarray(noisy)
        return image, target

class RandomRotation:
    def __init__(self,degrees,mean,ignore_value=255):
        self.degrees=degrees
        self.mean=mean
        self.ignore_value=ignore_value
    def __call__(self, image, target):
        expand=True
        if random.random()<0.5:
            angle = random.uniform(*self.degrees)
            image=F.rotate(image, angle,fill=self.mean,expand=expand)
            target=F.rotate(target,angle,fill=self.ignore_value,expand=expand)
        return image,target


class RandomScale(object):
    """
    Applies random scale augmentation.
    Arguments:
        min_scale: Minimum scale value.
        max_scale: Maximum scale value.
        scale_step_size: The step size from minimum to maximum value.
    """
    def __init__(self, min_scale, max_scale, scale_step_size):
        self.min_scale = min_scale
        self.max_scale = max_scale
        self.scale_step_size = scale_step_size

    @staticmethod
    def get_random_scale(min_scale_factor, max_scale_factor, step_size):
        """Gets a random scale value.
        Args:
            min_scale_factor: Minimum scale value.
            max_scale_factor: Maximum scale value.
            step_size: The step size from minimum to maximum value.
        Returns:
            A random scale value selected between minimum and maximum value.
        Raises:
            ValueError: min_scale_factor has unexpected value.
        """
        if min_scale_factor < 0 or min_scale_factor > max_scale_factor:
            raise ValueError('Unexpected value of min_scale_factor.')

        if min_scale_factor == max_scale_factor:
            return min_scale_factor

        # When step_size = 0, we sample the value uniformly from [min, max).
        if step_size == 0:
            return random.uniform(min_scale_factor, max_scale_factor)

        # When step_size != 0, we randomly select one discrete value from [min, max].
        num_steps = int((max_scale_factor - min_scale_factor) / step_size + 1)
        scale_factors = np.linspace(min_scale_factor, max_scale_factor, num_steps)
        np.random.shuffle(scale_factors)
        return scale_factors[0]

    def __call__(self, image, label):
        scale = self.get_random_scale(self.min_scale, self.max_scale, self.scale_step_size)
        img_w, img_h = image.size
        img_w,img_h=int(img_w*scale),int(img_h*scale)
        image=F.resize(image,[img_h,img_w])
        label=F.resize(label,[img_h,img_w],interpolation=F.InterpolationMode.NEAREST)
        return image,label


class RandomCrop(object):
    def __init__(self, crop_h, crop_w, pad_value, ignore_label, random_pad):
        self.crop_h = crop_h
        self.crop_w = crop_w
        self.pad_value = pad_value
        self.ignore_label = ignore_label
        self.random_pad = random_pad

    def __call__(self, image, label):
        img_w,img_h=image.size
        pad_h = max(self.crop_h - img_h, 0)
        pad_w = max(self.crop_w - img_w, 0)
        if pad_h > 0 or pad_w > 0:
            if self.random_pad:
                pad_top = random.randint(0, pad_h)
                pad_bottom = pad_h - pad_top
                pad_left = random.randint(0, pad_w)
                pad_right = pad_w - pad_left
            else:
                pad_top, pad_bottom, pad_left, pad_right = 0, pad_h, 0, pad_w
            image = F.pad(image, (pad_left, pad_top, pad_right, pad_bottom), fill=self.pad_value)
            label= F.pad(label, (pad_left, pad_top, pad_right, pad_bottom), fill=self.ignore_label)

        crop_params = T.RandomCrop.get_params(image, (self.crop_h, self.crop_w))
        image = F.crop(image, *crop_params)
        label = F.crop(label, *crop_params)
        return image,label


class RandomHorizontalFlip(object):
    def __init__(self, flip_prob):
        self.flip_prob = flip_prob

    def __call__(self, image, target):
        if random.random() < self.flip_prob:
            image = F.hflip(image)
            target = F.hflip(target)
        return image, target

def f():
    image=np.zeros((50,50),dtype=np.uint8)
    assert(image.dtype==np.uint8)
    gauss=(np.array(torch.randn(50,50)*10)).astype("uint8")
    noisy = (image + gauss).clip(0, 255)
    # print(noisy.dtype)
    # factor = random.uniform(0, 10)
    # image = Image.open("cityscapes_dataset/leftImg8bit/train/aachen/aachen_000000_000019_leftImg8bit.png")
    # image=np.array(image)
    # gauss = np.array(torch.randn(*image.shape)) * factor
    # noisy = (image + gauss).clip(0, 255).astype("uint8")

def g():
    image=np.zeros((50,50))
    gauss=np.array(torch.randn(50,50)*10)
    noisy = (image + gauss).clip(0, 255).astype("uint8")
    # print(noisy.dtype)
    # factor = random.uniform(0, 10)
    # image = Image.open("cityscapes_dataset/leftImg8bit/train/aachen/aachen_000000_000019_leftImg8bit.png")
    # image=np.array(image)
    # gauss = np.array(torch.randn(*image.shape)) * factor
    # print(gauss.dtype)
    # noisy = (image + gauss).clip(0, 255).astype("uint8")


if __name__=='__main__':
    f()
    g()
    #import timeit
    #print(timeit.timeit('f()', globals=globals(), number=100))
    #print(timeit.timeit('g()', globals=globals(), number=100))