| |
|
|
| """ |
| Appearance extractor(F) defined in paper, which maps the source image s to a 3D appearance feature volume. |
| """ |
|
|
| import torch |
| from torch import nn |
| from .util import SameBlock2d, DownBlock2d, ResBlock3d |
|
|
|
|
| class AppearanceFeatureExtractor(nn.Module): |
|
|
| def __init__(self, image_channel, block_expansion, num_down_blocks, max_features, reshape_channel, reshape_depth, num_resblocks): |
| super(AppearanceFeatureExtractor, self).__init__() |
| self.image_channel = image_channel |
| self.block_expansion = block_expansion |
| self.num_down_blocks = num_down_blocks |
| self.max_features = max_features |
| self.reshape_channel = reshape_channel |
| self.reshape_depth = reshape_depth |
|
|
| self.first = SameBlock2d(image_channel, block_expansion, kernel_size=(3, 3), padding=(1, 1)) |
|
|
| down_blocks = [] |
| for i in range(num_down_blocks): |
| in_features = min(max_features, block_expansion * (2 ** i)) |
| out_features = min(max_features, block_expansion * (2 ** (i + 1))) |
| down_blocks.append(DownBlock2d(in_features, out_features, kernel_size=(3, 3), padding=(1, 1))) |
| self.down_blocks = nn.ModuleList(down_blocks) |
|
|
| self.second = nn.Conv2d(in_channels=out_features, out_channels=max_features, kernel_size=1, stride=1) |
|
|
| self.resblocks_3d = torch.nn.Sequential() |
| for i in range(num_resblocks): |
| self.resblocks_3d.add_module('3dr' + str(i), ResBlock3d(reshape_channel, kernel_size=3, padding=1)) |
|
|
| def forward(self, source_image): |
| out = self.first(source_image) |
|
|
| for i in range(len(self.down_blocks)): |
| out = self.down_blocks[i](out) |
| out = self.second(out) |
| bs, c, h, w = out.shape |
|
|
| f_s = out.view(bs, self.reshape_channel, self.reshape_depth, h, w) |
| f_s = self.resblocks_3d(f_s) |
| return f_s |
|
|
