Spaces:
Runtime error
Runtime error
File size: 1,887 Bytes
d7db483 2651d59 d7db483 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
import math
import random
from PIL import Image
import blobfile as bf
#from mpi4py import MPI
import numpy as np
from torch.utils.data import DataLoader, Dataset
import os
import torchvision.transforms as transforms
import torch as th
from .degradation.bsrgan_light import degradation_bsrgan_variant as degradation_fn_bsr_light
from functools import partial
import cv2
def get_params( size, resize_size, crop_size):
w, h = size
new_h = h
new_w = w
ss, ls = min(w, h), max(w, h) # shortside and longside
width_is_shorter = w == ss
ls = int(resize_size * ls / ss)
ss = resize_size
new_w, new_h = (ss, ls) if width_is_shorter else (ls, ss)
x = random.randint(0, np.maximum(0, new_w - crop_size))
y = random.randint(0, np.maximum(0, new_h - crop_size))
flip = random.random() > 0.5
return {'crop_pos': (x, y), 'flip': flip}
def get_transform(params, resize_size, crop_size, method=Image.BICUBIC, flip=True, crop = True):
transform_list = []
transform_list.append(transforms.Lambda(lambda img: __scale(img, crop_size, method)))
if flip:
transform_list.append(transforms.Lambda(lambda img: __flip(img, params['flip'])))
return transforms.Compose(transform_list)
def get_tensor(normalize=True, toTensor=True):
transform_list = []
if toTensor:
transform_list += [transforms.ToTensor()]
if normalize:
transform_list += [transforms.Normalize((0.5, 0.5, 0.5),
(0.5, 0.5, 0.5))]
return transforms.Compose(transform_list)
def normalize():
return transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
def __scale(img, target_width, method=Image.BICUBIC):
return img.resize((target_width, target_width), method)
def __flip(img, flip):
if flip:
return img.transpose(Image.FLIP_LEFT_RIGHT)
return img |