Spaces:
Sleeping
Sleeping
File size: 986 Bytes
d8afa61 |
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 |
import cv2
import numpy as np
import torch
from torchvision import transforms
class Resize(object):
def __init__(self, size):
self.size = size
def __call__(self, image):
image = cv2.resize(image, (self.size, self.size))
return image
class NormalizeImage(object):
def __init__(self, mean, std):
self.mean = mean
self.std = std
def __call__(self, image):
image = image.astype(np.float32) / 255.0
image -= np.array(self.mean)
image /= np.array(self.std)
return image
class PrepareForNet(object):
def __call__(self, image):
image = torch.from_numpy(image)
if len(image.shape) == 3:
image = image.permute(2, 0, 1)
image = image.unsqueeze(0)
return image
class Compose:
def __init__(self, transforms):
self.transforms = transforms
def __call__(self, img):
for t in self.transforms:
img = t(img)
return img |