Shellbrady/LivePortrait5

1# coding: utf-82 3"""4Warping field estimator(W) defined in the paper, which generates a warping field using the implicit5keypoint representations x_s and x_d, and employs this flow field to warp the source feature volume f_s.6"""7 8from torch import nn9import torch.nn.functional as F10from .util import SameBlock2d11from .dense_motion import DenseMotionNetwork12 13 14class WarpingNetwork(nn.Module):15    def __init__(16        self,17        num_kp,18        block_expansion,19        max_features,20        num_down_blocks,21        reshape_channel,22        estimate_occlusion_map=False,23        dense_motion_params=None,24        **kwargs25    ):26        super(WarpingNetwork, self).__init__()27 28        self.upscale = kwargs.get('upscale', 1)29        self.flag_use_occlusion_map = kwargs.get('flag_use_occlusion_map', True)30 31        if dense_motion_params is not None:32            self.dense_motion_network = DenseMotionNetwork(33                num_kp=num_kp,34                feature_channel=reshape_channel,35                estimate_occlusion_map=estimate_occlusion_map,36                **dense_motion_params37            )38        else:39            self.dense_motion_network = None40 41        self.third = SameBlock2d(max_features, block_expansion * (2 ** num_down_blocks), kernel_size=(3, 3), padding=(1, 1), lrelu=True)42        self.fourth = nn.Conv2d(in_channels=block_expansion * (2 ** num_down_blocks), out_channels=block_expansion * (2 ** num_down_blocks), kernel_size=1, stride=1)43 44        self.estimate_occlusion_map = estimate_occlusion_map45 46    def deform_input(self, inp, deformation):47        return F.grid_sample(inp, deformation, align_corners=False)48 49    def forward(self, feature_3d, kp_driving, kp_source):50        if self.dense_motion_network is not None:51            # Feature warper, Transforming feature representation according to deformation and occlusion52            dense_motion = self.dense_motion_network(53                feature=feature_3d, kp_driving=kp_driving, kp_source=kp_source54            )55            if 'occlusion_map' in dense_motion:56                occlusion_map = dense_motion['occlusion_map']  # Bx1x64x6457            else:58                occlusion_map = None59 60            deformation = dense_motion['deformation']  # Bx16x64x64x361            out = self.deform_input(feature_3d, deformation)  # Bx32x16x64x6462 63            bs, c, d, h, w = out.shape  # Bx32x16x64x6464            out = out.view(bs, c * d, h, w)  # -> Bx512x64x6465            out = self.third(out)  # -> Bx256x64x6466            out = self.fourth(out)  # -> Bx256x64x6467 68            if self.flag_use_occlusion_map and (occlusion_map is not None):69                out = out * occlusion_map70 71        ret_dct = {72            'occlusion_map': occlusion_map,73            'deformation': deformation,74            'out': out,75        }76 77        return ret_dct78