wehub-resource-sync
58 lines
2.1 KiB
Python
58 lines
2.1 KiB
Python
# Copyright Howto100m authors.
|
|
# Copyright (c) Facebook, Inc. All Rights Reserved
|
|
|
|
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])
|
|
elif type == 'vmz':
|
|
self.norm = Normalize(mean=[110.201, 100.64, 95.997], std=[58.1489, 56.4701, 55.3324])
|
|
|
|
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 == 'vmz':
|
|
#tensor = self._zero_pad(tensor, 8)
|
|
tensor = self._zero_pad(tensor, 10)
|
|
tensor = self.norm(tensor)
|
|
#tensor = tensor.view(-1, 8, 3, 112, 112)
|
|
tensor = tensor.view(-1, 10, 3, 112, 112)
|
|
tensor = tensor.transpose(1, 2)
|
|
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)
|
|
elif self.type == 's3d':
|
|
tensor = tensor / 255.0
|
|
tensor = self._zero_pad(tensor, 30)
|
|
tensor = tensor.view(-1, 30, 3, 224, 224) # N x 30 x 3 x H x W
|
|
tensor = tensor.transpose(1, 2) # N x 3 x 30 x H x W
|
|
# for vae do nothing
|
|
return tensor
|