Spaces:
Runtime error
Runtime error
File size: 1,794 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 |
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 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 |