Spaces:
Running
on
T4
Running
on
T4
| import numpy as np | |
| import torch | |
| from torch import nn | |
| from torch.nn import functional as F | |
| from torchvision import transforms | |
| class ToOneHot(object): | |
| """ Convert the input PIL image to a one-hot torch tensor """ | |
| def __init__(self, n_classes=None): | |
| self.n_classes = n_classes | |
| def onehot_initialization(self, a): | |
| if self.n_classes is None: | |
| self.n_classes = len(np.unique(a)) | |
| out = np.zeros(a.shape + (self.n_classes, ), dtype=int) | |
| out[self.__all_idx(a, axis=2)] = 1 | |
| return out | |
| def __all_idx(self, idx, axis): | |
| grid = np.ogrid[tuple(map(slice, idx.shape))] | |
| grid.insert(axis, idx) | |
| return tuple(grid) | |
| def __call__(self, img): | |
| img = np.array(img) | |
| one_hot = self.onehot_initialization(img) | |
| return one_hot | |
| class BilinearResize(object): | |
| def __init__(self, factors=[1, 2, 4, 8, 16, 32]): | |
| self.factors = factors | |
| def __call__(self, image): | |
| factor = np.random.choice(self.factors, size=1)[0] | |
| D = BicubicDownSample(factor=factor, cuda=False) | |
| img_tensor = transforms.ToTensor()(image).unsqueeze(0) | |
| img_tensor_lr = D(img_tensor)[0].clamp(0, 1) | |
| img_low_res = transforms.ToPILImage()(img_tensor_lr) | |
| return img_low_res | |
| class BicubicDownSample(nn.Module): | |
| def bicubic_kernel(self, x, a=-0.50): | |
| """ | |
| This equation is exactly copied from the website below: | |
| https://clouard.users.greyc.fr/Pantheon/experiments/rescaling/index-en.html#bicubic | |
| """ | |
| abs_x = torch.abs(x) | |
| if abs_x <= 1.: | |
| return (a + 2.) * torch.pow(abs_x, 3.) - (a + 3.) * torch.pow(abs_x, 2.) + 1 | |
| elif 1. < abs_x < 2.: | |
| return a * torch.pow(abs_x, 3) - 5. * a * torch.pow(abs_x, 2.) + 8. * a * abs_x - 4. * a | |
| else: | |
| return 0.0 | |
| def __init__(self, factor=4, cuda=True, padding='reflect'): | |
| super().__init__() | |
| self.factor = factor | |
| size = factor * 4 | |
| k = torch.tensor([self.bicubic_kernel((i - torch.floor(torch.tensor(size / 2)) + 0.5) / factor) | |
| for i in range(size)], dtype=torch.float32) | |
| k = k / torch.sum(k) | |
| k1 = torch.reshape(k, shape=(1, 1, size, 1)) | |
| self.k1 = torch.cat([k1, k1, k1], dim=0) | |
| k2 = torch.reshape(k, shape=(1, 1, 1, size)) | |
| self.k2 = torch.cat([k2, k2, k2], dim=0) | |
| self.cuda = '.cuda' if cuda else '' | |
| self.padding = padding | |
| for param in self.parameters(): | |
| param.requires_grad = False | |
| def forward(self, x, nhwc=False, clip_round=False, byte_output=False): | |
| filter_height = self.factor * 4 | |
| filter_width = self.factor * 4 | |
| stride = self.factor | |
| pad_along_height = max(filter_height - stride, 0) | |
| pad_along_width = max(filter_width - stride, 0) | |
| filters1 = self.k1.type('torch{}.FloatTensor'.format(self.cuda)) | |
| filters2 = self.k2.type('torch{}.FloatTensor'.format(self.cuda)) | |
| # compute actual padding values for each side | |
| pad_top = pad_along_height // 2 | |
| pad_bottom = pad_along_height - pad_top | |
| pad_left = pad_along_width // 2 | |
| pad_right = pad_along_width - pad_left | |
| # apply mirror padding | |
| if nhwc: | |
| x = torch.transpose(torch.transpose(x, 2, 3), 1, 2) # NHWC to NCHW | |
| # downscaling performed by 1-d convolution | |
| x = F.pad(x, (0, 0, pad_top, pad_bottom), self.padding) | |
| x = F.conv2d(input=x, weight=filters1, stride=(stride, 1), groups=3) | |
| if clip_round: | |
| x = torch.clamp(torch.round(x), 0.0, 255.) | |
| x = F.pad(x, (pad_left, pad_right, 0, 0), self.padding) | |
| x = F.conv2d(input=x, weight=filters2, stride=(1, stride), groups=3) | |
| if clip_round: | |
| x = torch.clamp(torch.round(x), 0.0, 255.) | |
| if nhwc: | |
| x = torch.transpose(torch.transpose(x, 1, 3), 1, 2) | |
| if byte_output: | |
| return x.type('torch.ByteTensor'.format(self.cuda)) | |
| else: | |
| return x | |