Shellbrady/LivePortrait5

1# coding: utf-82 3"""4Appearance extractor(F) defined in paper, which maps the source image s to a 3D appearance feature volume.5"""6 7import torch8from torch import nn9from .util import SameBlock2d, DownBlock2d, ResBlock3d10 11 12class AppearanceFeatureExtractor(nn.Module):13 14    def __init__(self, image_channel, block_expansion, num_down_blocks, max_features, reshape_channel, reshape_depth, num_resblocks):15        super(AppearanceFeatureExtractor, self).__init__()16        self.image_channel = image_channel17        self.block_expansion = block_expansion18        self.num_down_blocks = num_down_blocks19        self.max_features = max_features20        self.reshape_channel = reshape_channel21        self.reshape_depth = reshape_depth22 23        self.first = SameBlock2d(image_channel, block_expansion, kernel_size=(3, 3), padding=(1, 1))24 25        down_blocks = []26        for i in range(num_down_blocks):27            in_features = min(max_features, block_expansion * (2 ** i))28            out_features = min(max_features, block_expansion * (2 ** (i + 1)))29            down_blocks.append(DownBlock2d(in_features, out_features, kernel_size=(3, 3), padding=(1, 1)))30        self.down_blocks = nn.ModuleList(down_blocks)31 32        self.second = nn.Conv2d(in_channels=out_features, out_channels=max_features, kernel_size=1, stride=1)33 34        self.resblocks_3d = torch.nn.Sequential()35        for i in range(num_resblocks):36            self.resblocks_3d.add_module('3dr' + str(i), ResBlock3d(reshape_channel, kernel_size=3, padding=1))37 38    def forward(self, source_image):39        out = self.first(source_image)  # Bx3x256x256 -> Bx64x256x25640 41        for i in range(len(self.down_blocks)):42            out = self.down_blocks[i](out)43        out = self.second(out)44        bs, c, h, w = out.shape  # ->Bx512x64x6445 46        f_s = out.view(bs, self.reshape_channel, self.reshape_depth, h, w)  # ->Bx32x16x64x6447        f_s = self.resblocks_3d(f_s)  # ->Bx32x16x64x6448        return f_s49