| |
| |
|
|
| 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, 10) |
| tensor = self.norm(tensor) |
| |
| 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) |
| tensor = tensor.transpose(1, 2) |
| |
| return tensor |
|
|
