Shellbrady/LivePortrait5

1# coding: utf-82 3"""4The module that predicting a dense motion from sparse motion representation given by kp_source and kp_driving5"""6 7from torch import nn8import torch.nn.functional as F9import torch10from .util import Hourglass, make_coordinate_grid, kp2gaussian11 12 13class DenseMotionNetwork(nn.Module):14    def __init__(self, block_expansion, num_blocks, max_features, num_kp, feature_channel, reshape_depth, compress, estimate_occlusion_map=True):15        super(DenseMotionNetwork, self).__init__()16        self.hourglass = Hourglass(block_expansion=block_expansion, in_features=(num_kp+1)*(compress+1), max_features=max_features, num_blocks=num_blocks)  # ~60+G17 18        self.mask = nn.Conv3d(self.hourglass.out_filters, num_kp + 1, kernel_size=7, padding=3)  # 65G! NOTE: computation cost is large19        self.compress = nn.Conv3d(feature_channel, compress, kernel_size=1)  # 0.8G20        self.norm = nn.BatchNorm3d(compress, affine=True)21        self.num_kp = num_kp22        self.flag_estimate_occlusion_map = estimate_occlusion_map23 24        if self.flag_estimate_occlusion_map:25            self.occlusion = nn.Conv2d(self.hourglass.out_filters*reshape_depth, 1, kernel_size=7, padding=3)26        else:27            self.occlusion = None28 29    def create_sparse_motions(self, feature, kp_driving, kp_source):30        bs, _, d, h, w = feature.shape  # (bs, 4, 16, 64, 64)31        identity_grid = make_coordinate_grid((d, h, w), ref=kp_source)  # (16, 64, 64, 3)32        identity_grid = identity_grid.view(1, 1, d, h, w, 3)  # (1, 1, d=16, h=64, w=64, 3)33        coordinate_grid = identity_grid - kp_driving.view(bs, self.num_kp, 1, 1, 1, 3)34 35        k = coordinate_grid.shape[1]36 37        # NOTE: there lacks an one-order flow38        driving_to_source = coordinate_grid + kp_source.view(bs, self.num_kp, 1, 1, 1, 3)    # (bs, num_kp, d, h, w, 3)39 40        # adding background feature41        identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1, 1)42        sparse_motions = torch.cat([identity_grid, driving_to_source], dim=1)  # (bs, 1+num_kp, d, h, w, 3)43        return sparse_motions44 45    def create_deformed_feature(self, feature, sparse_motions):46        bs, _, d, h, w = feature.shape47        feature_repeat = feature.unsqueeze(1).unsqueeze(1).repeat(1, self.num_kp+1, 1, 1, 1, 1, 1)      # (bs, num_kp+1, 1, c, d, h, w)48        feature_repeat = feature_repeat.view(bs * (self.num_kp+1), -1, d, h, w)                         # (bs*(num_kp+1), c, d, h, w)49        sparse_motions = sparse_motions.view((bs * (self.num_kp+1), d, h, w, -1))                       # (bs*(num_kp+1), d, h, w, 3)50        sparse_deformed = F.grid_sample(feature_repeat, sparse_motions, align_corners=False)51        sparse_deformed = sparse_deformed.view((bs, self.num_kp+1, -1, d, h, w))                        # (bs, num_kp+1, c, d, h, w)52 53        return sparse_deformed54 55    def create_heatmap_representations(self, feature, kp_driving, kp_source):56        spatial_size = feature.shape[3:]  # (d=16, h=64, w=64)57        gaussian_driving = kp2gaussian(kp_driving, spatial_size=spatial_size, kp_variance=0.01)  # (bs, num_kp, d, h, w)58        gaussian_source = kp2gaussian(kp_source, spatial_size=spatial_size, kp_variance=0.01)  # (bs, num_kp, d, h, w)59        heatmap = gaussian_driving - gaussian_source  # (bs, num_kp, d, h, w)60 61        # adding background feature62        zeros = torch.zeros(heatmap.shape[0], 1, spatial_size[0], spatial_size[1], spatial_size[2]).type(heatmap.type()).to(heatmap.device)63        heatmap = torch.cat([zeros, heatmap], dim=1)64        heatmap = heatmap.unsqueeze(2)         # (bs, 1+num_kp, 1, d, h, w)65        return heatmap66 67    def forward(self, feature, kp_driving, kp_source):68        bs, _, d, h, w = feature.shape  # (bs, 32, 16, 64, 64)69 70        feature = self.compress(feature)  # (bs, 4, 16, 64, 64)71        feature = self.norm(feature)  # (bs, 4, 16, 64, 64)72        feature = F.relu(feature)  # (bs, 4, 16, 64, 64)73 74        out_dict = dict()75 76        # 1. deform 3d feature77        sparse_motion = self.create_sparse_motions(feature, kp_driving, kp_source)  # (bs, 1+num_kp, d, h, w, 3)78        deformed_feature = self.create_deformed_feature(feature, sparse_motion)  # (bs, 1+num_kp, c=4, d=16, h=64, w=64)79 80        # 2. (bs, 1+num_kp, d, h, w)81        heatmap = self.create_heatmap_representations(deformed_feature, kp_driving, kp_source)  # (bs, 1+num_kp, 1, d, h, w)82 83        input = torch.cat([heatmap, deformed_feature], dim=2)  # (bs, 1+num_kp, c=5, d=16, h=64, w=64)84        input = input.view(bs, -1, d, h, w)  # (bs, (1+num_kp)*c=105, d=16, h=64, w=64)85 86        prediction = self.hourglass(input)87 88        mask = self.mask(prediction)89        mask = F.softmax(mask, dim=1)  # (bs, 1+num_kp, d=16, h=64, w=64)90        out_dict['mask'] = mask91        mask = mask.unsqueeze(2)                                   # (bs, num_kp+1, 1, d, h, w)92        sparse_motion = sparse_motion.permute(0, 1, 5, 2, 3, 4)    # (bs, num_kp+1, 3, d, h, w)93        deformation = (sparse_motion * mask).sum(dim=1)            # (bs, 3, d, h, w)  mask take effect in this place94        deformation = deformation.permute(0, 2, 3, 4, 1)           # (bs, d, h, w, 3)95 96        out_dict['deformation'] = deformation97 98        if self.flag_estimate_occlusion_map:99            bs, _, d, h, w = prediction.shape100            prediction_reshape = prediction.view(bs, -1, h, w)101            occlusion_map = torch.sigmoid(self.occlusion(prediction_reshape))  # Bx1x64x64102            out_dict['occlusion_map'] = occlusion_map103 104        return out_dict105