Spaces:
Sleeping
Sleeping
File size: 1,255 Bytes
4e46a55 |
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 torch as th
class Normalize(object):
def __init__(self, mean, std):
self.mean = th.FloatTensor(mean).view(1, 3, 1, 1)
self.std = th.FloatTensor(std).view(1, 3, 1, 1)
def __call__(self, tensor):
tensor = (tensor - self.mean) / (self.std + 1e-8)
return tensor
class Preprocessing(object):
def __init__(self, type):
self.type = type
if type == '2d':
self.norm = Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
elif type == '3d':
self.norm = Normalize(mean=[110.6, 103.2, 96.3], std=[1.0, 1.0, 1.0])
def _zero_pad(self, tensor, size):
n = size - len(tensor) % size
if n == size:
return tensor
else:
z = th.zeros(n, tensor.shape[1], tensor.shape[2], tensor.shape[3])
return th.cat((tensor, z), 0)
def __call__(self, tensor):
if self.type == '2d':
tensor = tensor / 255.0
tensor = self.norm(tensor)
elif self.type == '3d':
tensor = self._zero_pad(tensor, 16)
tensor = self.norm(tensor)
tensor = tensor.view(-1, 16, 3, 112, 112)
tensor = tensor.transpose(1, 2)
return tensor
|