import random

import numpy as np
import torch
from PIL import Image, ImageFilter


class RandomVerticalFlip(object):
    def __call__(self, img):
        if random.random() < 0.5:
            return img.transpose(Image.FLIP_TOP_BOTTOM)
        return img


class RandomGaussianBlur(object):
    def __call__(self, img):
        if random.random() < 0.2:
            return img.filter(ImageFilter.BLUR)
        return img


class DeNormalize(object):
    def __init__(self, mean, std):
        self.mean = mean
        self.std = std

    def __call__(self, tensor):
        for t, m, s in zip(tensor, self.mean, self.std):
            t.mul_(s).add_(m)
        return tensor


class MaskToTensor(object):
    def __call__(self, img):
        return torch.from_numpy(np.array(img, dtype=np.int32)).long()


class FreeScale(object):
    def __init__(self, size, interpolation=Image.BILINEAR):
        self.size = size  # (w, h)
        self.interpolation = interpolation

    def __call__(self, img):
        return img.resize(self.size, self.interpolation)


class ChangeLabel(object):
    def __init__(self, ori_label, new_label):
        self.ori_label = ori_label
        self.new_label = new_label

    def __call__(self, mask):
        mask[mask == self.ori_label] = self.new_label
        return mask
