Spaces:
Sleeping
Sleeping
| # -*- coding: utf-8 -*- | |
| # | |
| # Developed by Haozhe Xie <cshzxie@gmail.com> | |
| import torch | |
| class Refiner(torch.nn.Module): | |
| def __init__(self, cfg): | |
| super(Refiner, self).__init__() | |
| self.cfg = cfg | |
| # Layer Definition | |
| self.layer1 = torch.nn.Sequential( | |
| torch.nn.Conv3d(1, 32, kernel_size=4, padding=2), | |
| torch.nn.BatchNorm3d(32), | |
| torch.nn.LeakyReLU(cfg.NETWORK.LEAKY_VALUE), | |
| torch.nn.MaxPool3d(kernel_size=2) | |
| ) | |
| self.layer2 = torch.nn.Sequential( | |
| torch.nn.Conv3d(32, 64, kernel_size=4, padding=2), | |
| torch.nn.BatchNorm3d(64), | |
| torch.nn.LeakyReLU(cfg.NETWORK.LEAKY_VALUE), | |
| torch.nn.MaxPool3d(kernel_size=2) | |
| ) | |
| self.layer3 = torch.nn.Sequential( | |
| torch.nn.Conv3d(64, 128, kernel_size=4, padding=2), | |
| torch.nn.BatchNorm3d(128), | |
| torch.nn.LeakyReLU(cfg.NETWORK.LEAKY_VALUE), | |
| torch.nn.MaxPool3d(kernel_size=2) | |
| ) | |
| self.layer4 = torch.nn.Sequential( | |
| torch.nn.Linear(8192, 2048), | |
| torch.nn.ReLU() | |
| ) | |
| self.layer5 = torch.nn.Sequential( | |
| torch.nn.Linear(2048, 8192), | |
| torch.nn.ReLU() | |
| ) | |
| self.layer6 = torch.nn.Sequential( | |
| torch.nn.ConvTranspose3d(128, 64, kernel_size=4, stride=2, bias=cfg.NETWORK.TCONV_USE_BIAS, padding=1), | |
| torch.nn.BatchNorm3d(64), | |
| torch.nn.ReLU() | |
| ) | |
| self.layer7 = torch.nn.Sequential( | |
| torch.nn.ConvTranspose3d(64, 32, kernel_size=4, stride=2, bias=cfg.NETWORK.TCONV_USE_BIAS, padding=1), | |
| torch.nn.BatchNorm3d(32), | |
| torch.nn.ReLU() | |
| ) | |
| self.layer8 = torch.nn.Sequential( | |
| torch.nn.ConvTranspose3d(32, 1, kernel_size=4, stride=2, bias=cfg.NETWORK.TCONV_USE_BIAS, padding=1), | |
| torch.nn.Sigmoid() | |
| ) | |
| def forward(self, coarse_volumes): | |
| volumes_32_l = coarse_volumes.view((-1, 1, self.cfg.CONST.N_VOX, self.cfg.CONST.N_VOX, self.cfg.CONST.N_VOX)) | |
| # print(volumes_32_l.size()) # torch.Size([batch_size, 1, 32, 32, 32]) | |
| volumes_16_l = self.layer1(volumes_32_l) | |
| # print(volumes_16_l.size()) # torch.Size([batch_size, 32, 16, 16, 16]) | |
| volumes_8_l = self.layer2(volumes_16_l) | |
| # print(volumes_8_l.size()) # torch.Size([batch_size, 64, 8, 8, 8]) | |
| volumes_4_l = self.layer3(volumes_8_l) | |
| # print(volumes_4_l.size()) # torch.Size([batch_size, 128, 4, 4, 4]) | |
| flatten_features = self.layer4(volumes_4_l.view(-1, 8192)) | |
| # print(flatten_features.size()) # torch.Size([batch_size, 2048]) | |
| flatten_features = self.layer5(flatten_features) | |
| # print(flatten_features.size()) # torch.Size([batch_size, 8192]) | |
| volumes_4_r = volumes_4_l + flatten_features.view(-1, 128, 4, 4, 4) | |
| # print(volumes_4_r.size()) # torch.Size([batch_size, 128, 4, 4, 4]) | |
| volumes_8_r = volumes_8_l + self.layer6(volumes_4_r) | |
| # print(volumes_8_r.size()) # torch.Size([batch_size, 64, 8, 8, 8]) | |
| volumes_16_r = volumes_16_l + self.layer7(volumes_8_r) | |
| # print(volumes_16_r.size()) # torch.Size([batch_size, 32, 16, 16, 16]) | |
| volumes_32_r = (volumes_32_l + self.layer8(volumes_16_r)) * 0.5 | |
| # print(volumes_32_r.size()) # torch.Size([batch_size, 1, 32, 32, 32]) | |
| return volumes_32_r.view((-1, self.cfg.CONST.N_VOX, self.cfg.CONST.N_VOX, self.cfg.CONST.N_VOX)) | |