mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-17 00:23:14 +00:00
Add SadTalker @ ae0b26359d44c9379450a8e02a38b45eeec15549
MIT license Commit: https://github.com/OpenTalker/SadTalker commit ae0b26359d44c9379450a8e02a38b45eeec15549 (HEAD -> main, origin/main, origin/HEAD) Author: shadowcun <shadowcun@tencent.com> Date: Mon Jul 3 13:48:35 2023 +0800 fixed bug: hacking facexlib of FAN
This commit is contained in:
Brandon Thomas
@@ -0,0 +1,121 @@
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
import torch
|
||||
from src.facerender.modules.util import Hourglass, make_coordinate_grid, kp2gaussian
|
||||
|
||||
from src.facerender.sync_batchnorm import SynchronizedBatchNorm3d as BatchNorm3d
|
||||
|
||||
|
||||
class DenseMotionNetwork(nn.Module):
|
||||
"""
|
||||
Module that predicting a dense motion from sparse motion representation given by kp_source and kp_driving
|
||||
"""
|
||||
|
||||
def __init__(self, block_expansion, num_blocks, max_features, num_kp, feature_channel, reshape_depth, compress,
|
||||
estimate_occlusion_map=False):
|
||||
super(DenseMotionNetwork, self).__init__()
|
||||
# self.hourglass = Hourglass(block_expansion=block_expansion, in_features=(num_kp+1)*(feature_channel+1), max_features=max_features, num_blocks=num_blocks)
|
||||
self.hourglass = Hourglass(block_expansion=block_expansion, in_features=(num_kp+1)*(compress+1), max_features=max_features, num_blocks=num_blocks)
|
||||
|
||||
self.mask = nn.Conv3d(self.hourglass.out_filters, num_kp + 1, kernel_size=7, padding=3)
|
||||
|
||||
self.compress = nn.Conv3d(feature_channel, compress, kernel_size=1)
|
||||
self.norm = BatchNorm3d(compress, affine=True)
|
||||
|
||||
if estimate_occlusion_map:
|
||||
# self.occlusion = nn.Conv2d(reshape_channel*reshape_depth, 1, kernel_size=7, padding=3)
|
||||
self.occlusion = nn.Conv2d(self.hourglass.out_filters*reshape_depth, 1, kernel_size=7, padding=3)
|
||||
else:
|
||||
self.occlusion = None
|
||||
|
||||
self.num_kp = num_kp
|
||||
|
||||
|
||||
def create_sparse_motions(self, feature, kp_driving, kp_source):
|
||||
bs, _, d, h, w = feature.shape
|
||||
identity_grid = make_coordinate_grid((d, h, w), type=kp_source['value'].type())
|
||||
identity_grid = identity_grid.view(1, 1, d, h, w, 3)
|
||||
coordinate_grid = identity_grid - kp_driving['value'].view(bs, self.num_kp, 1, 1, 1, 3)
|
||||
|
||||
# if 'jacobian' in kp_driving:
|
||||
if 'jacobian' in kp_driving and kp_driving['jacobian'] is not None:
|
||||
jacobian = torch.matmul(kp_source['jacobian'], torch.inverse(kp_driving['jacobian']))
|
||||
jacobian = jacobian.unsqueeze(-3).unsqueeze(-3).unsqueeze(-3)
|
||||
jacobian = jacobian.repeat(1, 1, d, h, w, 1, 1)
|
||||
coordinate_grid = torch.matmul(jacobian, coordinate_grid.unsqueeze(-1))
|
||||
coordinate_grid = coordinate_grid.squeeze(-1)
|
||||
|
||||
|
||||
driving_to_source = coordinate_grid + kp_source['value'].view(bs, self.num_kp, 1, 1, 1, 3) # (bs, num_kp, d, h, w, 3)
|
||||
|
||||
#adding background feature
|
||||
identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1, 1)
|
||||
sparse_motions = torch.cat([identity_grid, driving_to_source], dim=1) #bs num_kp+1 d h w 3
|
||||
|
||||
# sparse_motions = driving_to_source
|
||||
|
||||
return sparse_motions
|
||||
|
||||
def create_deformed_feature(self, feature, sparse_motions):
|
||||
bs, _, d, h, w = feature.shape
|
||||
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)
|
||||
feature_repeat = feature_repeat.view(bs * (self.num_kp+1), -1, d, h, w) # (bs*(num_kp+1), c, d, h, w)
|
||||
sparse_motions = sparse_motions.view((bs * (self.num_kp+1), d, h, w, -1)) # (bs*(num_kp+1), d, h, w, 3) !!!!
|
||||
sparse_deformed = F.grid_sample(feature_repeat, sparse_motions)
|
||||
sparse_deformed = sparse_deformed.view((bs, self.num_kp+1, -1, d, h, w)) # (bs, num_kp+1, c, d, h, w)
|
||||
return sparse_deformed
|
||||
|
||||
def create_heatmap_representations(self, feature, kp_driving, kp_source):
|
||||
spatial_size = feature.shape[3:]
|
||||
gaussian_driving = kp2gaussian(kp_driving, spatial_size=spatial_size, kp_variance=0.01)
|
||||
gaussian_source = kp2gaussian(kp_source, spatial_size=spatial_size, kp_variance=0.01)
|
||||
heatmap = gaussian_driving - gaussian_source
|
||||
|
||||
# adding background feature
|
||||
zeros = torch.zeros(heatmap.shape[0], 1, spatial_size[0], spatial_size[1], spatial_size[2]).type(heatmap.type())
|
||||
heatmap = torch.cat([zeros, heatmap], dim=1)
|
||||
heatmap = heatmap.unsqueeze(2) # (bs, num_kp+1, 1, d, h, w)
|
||||
return heatmap
|
||||
|
||||
def forward(self, feature, kp_driving, kp_source):
|
||||
bs, _, d, h, w = feature.shape
|
||||
|
||||
feature = self.compress(feature)
|
||||
feature = self.norm(feature)
|
||||
feature = F.relu(feature)
|
||||
|
||||
out_dict = dict()
|
||||
sparse_motion = self.create_sparse_motions(feature, kp_driving, kp_source)
|
||||
deformed_feature = self.create_deformed_feature(feature, sparse_motion)
|
||||
|
||||
heatmap = self.create_heatmap_representations(deformed_feature, kp_driving, kp_source)
|
||||
|
||||
input_ = torch.cat([heatmap, deformed_feature], dim=2)
|
||||
input_ = input_.view(bs, -1, d, h, w)
|
||||
|
||||
# input = deformed_feature.view(bs, -1, d, h, w) # (bs, num_kp+1 * c, d, h, w)
|
||||
|
||||
prediction = self.hourglass(input_)
|
||||
|
||||
|
||||
mask = self.mask(prediction)
|
||||
mask = F.softmax(mask, dim=1)
|
||||
out_dict['mask'] = mask
|
||||
mask = mask.unsqueeze(2) # (bs, num_kp+1, 1, d, h, w)
|
||||
|
||||
zeros_mask = torch.zeros_like(mask)
|
||||
mask = torch.where(mask < 1e-3, zeros_mask, mask)
|
||||
|
||||
sparse_motion = sparse_motion.permute(0, 1, 5, 2, 3, 4) # (bs, num_kp+1, 3, d, h, w)
|
||||
deformation = (sparse_motion * mask).sum(dim=1) # (bs, 3, d, h, w)
|
||||
deformation = deformation.permute(0, 2, 3, 4, 1) # (bs, d, h, w, 3)
|
||||
|
||||
out_dict['deformation'] = deformation
|
||||
|
||||
if self.occlusion:
|
||||
bs, c, d, h, w = prediction.shape
|
||||
prediction = prediction.view(bs, -1, h, w)
|
||||
occlusion_map = torch.sigmoid(self.occlusion(prediction))
|
||||
out_dict['occlusion_map'] = occlusion_map
|
||||
|
||||
return out_dict
|
||||
@@ -0,0 +1,90 @@
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
from facerender.modules.util import kp2gaussian
|
||||
import torch
|
||||
|
||||
|
||||
class DownBlock2d(nn.Module):
|
||||
"""
|
||||
Simple block for processing video (encoder).
|
||||
"""
|
||||
|
||||
def __init__(self, in_features, out_features, norm=False, kernel_size=4, pool=False, sn=False):
|
||||
super(DownBlock2d, self).__init__()
|
||||
self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size)
|
||||
|
||||
if sn:
|
||||
self.conv = nn.utils.spectral_norm(self.conv)
|
||||
|
||||
if norm:
|
||||
self.norm = nn.InstanceNorm2d(out_features, affine=True)
|
||||
else:
|
||||
self.norm = None
|
||||
self.pool = pool
|
||||
|
||||
def forward(self, x):
|
||||
out = x
|
||||
out = self.conv(out)
|
||||
if self.norm:
|
||||
out = self.norm(out)
|
||||
out = F.leaky_relu(out, 0.2)
|
||||
if self.pool:
|
||||
out = F.avg_pool2d(out, (2, 2))
|
||||
return out
|
||||
|
||||
|
||||
class Discriminator(nn.Module):
|
||||
"""
|
||||
Discriminator similar to Pix2Pix
|
||||
"""
|
||||
|
||||
def __init__(self, num_channels=3, block_expansion=64, num_blocks=4, max_features=512,
|
||||
sn=False, **kwargs):
|
||||
super(Discriminator, self).__init__()
|
||||
|
||||
down_blocks = []
|
||||
for i in range(num_blocks):
|
||||
down_blocks.append(
|
||||
DownBlock2d(num_channels if i == 0 else min(max_features, block_expansion * (2 ** i)),
|
||||
min(max_features, block_expansion * (2 ** (i + 1))),
|
||||
norm=(i != 0), kernel_size=4, pool=(i != num_blocks - 1), sn=sn))
|
||||
|
||||
self.down_blocks = nn.ModuleList(down_blocks)
|
||||
self.conv = nn.Conv2d(self.down_blocks[-1].conv.out_channels, out_channels=1, kernel_size=1)
|
||||
if sn:
|
||||
self.conv = nn.utils.spectral_norm(self.conv)
|
||||
|
||||
def forward(self, x):
|
||||
feature_maps = []
|
||||
out = x
|
||||
|
||||
for down_block in self.down_blocks:
|
||||
feature_maps.append(down_block(out))
|
||||
out = feature_maps[-1]
|
||||
prediction_map = self.conv(out)
|
||||
|
||||
return feature_maps, prediction_map
|
||||
|
||||
|
||||
class MultiScaleDiscriminator(nn.Module):
|
||||
"""
|
||||
Multi-scale (scale) discriminator
|
||||
"""
|
||||
|
||||
def __init__(self, scales=(), **kwargs):
|
||||
super(MultiScaleDiscriminator, self).__init__()
|
||||
self.scales = scales
|
||||
discs = {}
|
||||
for scale in scales:
|
||||
discs[str(scale).replace('.', '-')] = Discriminator(**kwargs)
|
||||
self.discs = nn.ModuleDict(discs)
|
||||
|
||||
def forward(self, x):
|
||||
out_dict = {}
|
||||
for scale, disc in self.discs.items():
|
||||
scale = str(scale).replace('-', '.')
|
||||
key = 'prediction_' + scale
|
||||
feature_maps, prediction_map = disc(x[key])
|
||||
out_dict['feature_maps_' + scale] = feature_maps
|
||||
out_dict['prediction_map_' + scale] = prediction_map
|
||||
return out_dict
|
||||
@@ -0,0 +1,255 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
from src.facerender.modules.util import ResBlock2d, SameBlock2d, UpBlock2d, DownBlock2d, ResBlock3d, SPADEResnetBlock
|
||||
from src.facerender.modules.dense_motion import DenseMotionNetwork
|
||||
|
||||
|
||||
class OcclusionAwareGenerator(nn.Module):
|
||||
"""
|
||||
Generator follows NVIDIA architecture.
|
||||
"""
|
||||
|
||||
def __init__(self, image_channel, feature_channel, num_kp, block_expansion, max_features, num_down_blocks, reshape_channel, reshape_depth,
|
||||
num_resblocks, estimate_occlusion_map=False, dense_motion_params=None, estimate_jacobian=False):
|
||||
super(OcclusionAwareGenerator, self).__init__()
|
||||
|
||||
if dense_motion_params is not None:
|
||||
self.dense_motion_network = DenseMotionNetwork(num_kp=num_kp, feature_channel=feature_channel,
|
||||
estimate_occlusion_map=estimate_occlusion_map,
|
||||
**dense_motion_params)
|
||||
else:
|
||||
self.dense_motion_network = None
|
||||
|
||||
self.first = SameBlock2d(image_channel, block_expansion, kernel_size=(7, 7), padding=(3, 3))
|
||||
|
||||
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.reshape_channel = reshape_channel
|
||||
self.reshape_depth = reshape_depth
|
||||
|
||||
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))
|
||||
|
||||
out_features = block_expansion * (2 ** (num_down_blocks))
|
||||
self.third = SameBlock2d(max_features, out_features, kernel_size=(3, 3), padding=(1, 1), lrelu=True)
|
||||
self.fourth = nn.Conv2d(in_channels=out_features, out_channels=out_features, kernel_size=1, stride=1)
|
||||
|
||||
self.resblocks_2d = torch.nn.Sequential()
|
||||
for i in range(num_resblocks):
|
||||
self.resblocks_2d.add_module('2dr' + str(i), ResBlock2d(out_features, kernel_size=3, padding=1))
|
||||
|
||||
up_blocks = []
|
||||
for i in range(num_down_blocks):
|
||||
in_features = max(block_expansion, block_expansion * (2 ** (num_down_blocks - i)))
|
||||
out_features = max(block_expansion, block_expansion * (2 ** (num_down_blocks - i - 1)))
|
||||
up_blocks.append(UpBlock2d(in_features, out_features, kernel_size=(3, 3), padding=(1, 1)))
|
||||
self.up_blocks = nn.ModuleList(up_blocks)
|
||||
|
||||
self.final = nn.Conv2d(block_expansion, image_channel, kernel_size=(7, 7), padding=(3, 3))
|
||||
self.estimate_occlusion_map = estimate_occlusion_map
|
||||
self.image_channel = image_channel
|
||||
|
||||
def deform_input(self, inp, deformation):
|
||||
_, d_old, h_old, w_old, _ = deformation.shape
|
||||
_, _, d, h, w = inp.shape
|
||||
if d_old != d or h_old != h or w_old != w:
|
||||
deformation = deformation.permute(0, 4, 1, 2, 3)
|
||||
deformation = F.interpolate(deformation, size=(d, h, w), mode='trilinear')
|
||||
deformation = deformation.permute(0, 2, 3, 4, 1)
|
||||
return F.grid_sample(inp, deformation)
|
||||
|
||||
def forward(self, source_image, kp_driving, kp_source):
|
||||
# Encoding (downsampling) part
|
||||
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
|
||||
# print(out.shape)
|
||||
feature_3d = out.view(bs, self.reshape_channel, self.reshape_depth, h ,w)
|
||||
feature_3d = self.resblocks_3d(feature_3d)
|
||||
|
||||
# Transforming feature representation according to deformation and occlusion
|
||||
output_dict = {}
|
||||
if self.dense_motion_network is not None:
|
||||
dense_motion = self.dense_motion_network(feature=feature_3d, kp_driving=kp_driving,
|
||||
kp_source=kp_source)
|
||||
output_dict['mask'] = dense_motion['mask']
|
||||
|
||||
if 'occlusion_map' in dense_motion:
|
||||
occlusion_map = dense_motion['occlusion_map']
|
||||
output_dict['occlusion_map'] = occlusion_map
|
||||
else:
|
||||
occlusion_map = None
|
||||
deformation = dense_motion['deformation']
|
||||
out = self.deform_input(feature_3d, deformation)
|
||||
|
||||
bs, c, d, h, w = out.shape
|
||||
out = out.view(bs, c*d, h, w)
|
||||
out = self.third(out)
|
||||
out = self.fourth(out)
|
||||
|
||||
if occlusion_map is not None:
|
||||
if out.shape[2] != occlusion_map.shape[2] or out.shape[3] != occlusion_map.shape[3]:
|
||||
occlusion_map = F.interpolate(occlusion_map, size=out.shape[2:], mode='bilinear')
|
||||
out = out * occlusion_map
|
||||
|
||||
# output_dict["deformed"] = self.deform_input(source_image, deformation) # 3d deformation cannot deform 2d image
|
||||
|
||||
# Decoding part
|
||||
out = self.resblocks_2d(out)
|
||||
for i in range(len(self.up_blocks)):
|
||||
out = self.up_blocks[i](out)
|
||||
out = self.final(out)
|
||||
out = F.sigmoid(out)
|
||||
|
||||
output_dict["prediction"] = out
|
||||
|
||||
return output_dict
|
||||
|
||||
|
||||
class SPADEDecoder(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
ic = 256
|
||||
oc = 64
|
||||
norm_G = 'spadespectralinstance'
|
||||
label_nc = 256
|
||||
|
||||
self.fc = nn.Conv2d(ic, 2 * ic, 3, padding=1)
|
||||
self.G_middle_0 = SPADEResnetBlock(2 * ic, 2 * ic, norm_G, label_nc)
|
||||
self.G_middle_1 = SPADEResnetBlock(2 * ic, 2 * ic, norm_G, label_nc)
|
||||
self.G_middle_2 = SPADEResnetBlock(2 * ic, 2 * ic, norm_G, label_nc)
|
||||
self.G_middle_3 = SPADEResnetBlock(2 * ic, 2 * ic, norm_G, label_nc)
|
||||
self.G_middle_4 = SPADEResnetBlock(2 * ic, 2 * ic, norm_G, label_nc)
|
||||
self.G_middle_5 = SPADEResnetBlock(2 * ic, 2 * ic, norm_G, label_nc)
|
||||
self.up_0 = SPADEResnetBlock(2 * ic, ic, norm_G, label_nc)
|
||||
self.up_1 = SPADEResnetBlock(ic, oc, norm_G, label_nc)
|
||||
self.conv_img = nn.Conv2d(oc, 3, 3, padding=1)
|
||||
self.up = nn.Upsample(scale_factor=2)
|
||||
|
||||
def forward(self, feature):
|
||||
seg = feature
|
||||
x = self.fc(feature)
|
||||
x = self.G_middle_0(x, seg)
|
||||
x = self.G_middle_1(x, seg)
|
||||
x = self.G_middle_2(x, seg)
|
||||
x = self.G_middle_3(x, seg)
|
||||
x = self.G_middle_4(x, seg)
|
||||
x = self.G_middle_5(x, seg)
|
||||
x = self.up(x)
|
||||
x = self.up_0(x, seg) # 256, 128, 128
|
||||
x = self.up(x)
|
||||
x = self.up_1(x, seg) # 64, 256, 256
|
||||
|
||||
x = self.conv_img(F.leaky_relu(x, 2e-1))
|
||||
# x = torch.tanh(x)
|
||||
x = F.sigmoid(x)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class OcclusionAwareSPADEGenerator(nn.Module):
|
||||
|
||||
def __init__(self, image_channel, feature_channel, num_kp, block_expansion, max_features, num_down_blocks, reshape_channel, reshape_depth,
|
||||
num_resblocks, estimate_occlusion_map=False, dense_motion_params=None, estimate_jacobian=False):
|
||||
super(OcclusionAwareSPADEGenerator, self).__init__()
|
||||
|
||||
if dense_motion_params is not None:
|
||||
self.dense_motion_network = DenseMotionNetwork(num_kp=num_kp, feature_channel=feature_channel,
|
||||
estimate_occlusion_map=estimate_occlusion_map,
|
||||
**dense_motion_params)
|
||||
else:
|
||||
self.dense_motion_network = None
|
||||
|
||||
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.reshape_channel = reshape_channel
|
||||
self.reshape_depth = reshape_depth
|
||||
|
||||
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))
|
||||
|
||||
out_features = block_expansion * (2 ** (num_down_blocks))
|
||||
self.third = SameBlock2d(max_features, out_features, kernel_size=(3, 3), padding=(1, 1), lrelu=True)
|
||||
self.fourth = nn.Conv2d(in_channels=out_features, out_channels=out_features, kernel_size=1, stride=1)
|
||||
|
||||
self.estimate_occlusion_map = estimate_occlusion_map
|
||||
self.image_channel = image_channel
|
||||
|
||||
self.decoder = SPADEDecoder()
|
||||
|
||||
def deform_input(self, inp, deformation):
|
||||
_, d_old, h_old, w_old, _ = deformation.shape
|
||||
_, _, d, h, w = inp.shape
|
||||
if d_old != d or h_old != h or w_old != w:
|
||||
deformation = deformation.permute(0, 4, 1, 2, 3)
|
||||
deformation = F.interpolate(deformation, size=(d, h, w), mode='trilinear')
|
||||
deformation = deformation.permute(0, 2, 3, 4, 1)
|
||||
return F.grid_sample(inp, deformation)
|
||||
|
||||
def forward(self, source_image, kp_driving, kp_source):
|
||||
# Encoding (downsampling) part
|
||||
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
|
||||
# print(out.shape)
|
||||
feature_3d = out.view(bs, self.reshape_channel, self.reshape_depth, h ,w)
|
||||
feature_3d = self.resblocks_3d(feature_3d)
|
||||
|
||||
# Transforming feature representation according to deformation and occlusion
|
||||
output_dict = {}
|
||||
if self.dense_motion_network is not None:
|
||||
dense_motion = self.dense_motion_network(feature=feature_3d, kp_driving=kp_driving,
|
||||
kp_source=kp_source)
|
||||
output_dict['mask'] = dense_motion['mask']
|
||||
|
||||
# import pdb; pdb.set_trace()
|
||||
|
||||
if 'occlusion_map' in dense_motion:
|
||||
occlusion_map = dense_motion['occlusion_map']
|
||||
output_dict['occlusion_map'] = occlusion_map
|
||||
else:
|
||||
occlusion_map = None
|
||||
deformation = dense_motion['deformation']
|
||||
out = self.deform_input(feature_3d, deformation)
|
||||
|
||||
bs, c, d, h, w = out.shape
|
||||
out = out.view(bs, c*d, h, w)
|
||||
out = self.third(out)
|
||||
out = self.fourth(out)
|
||||
|
||||
# occlusion_map = torch.where(occlusion_map < 0.95, 0, occlusion_map)
|
||||
|
||||
if occlusion_map is not None:
|
||||
if out.shape[2] != occlusion_map.shape[2] or out.shape[3] != occlusion_map.shape[3]:
|
||||
occlusion_map = F.interpolate(occlusion_map, size=out.shape[2:], mode='bilinear')
|
||||
out = out * occlusion_map
|
||||
|
||||
# Decoding part
|
||||
out = self.decoder(out)
|
||||
|
||||
output_dict["prediction"] = out
|
||||
|
||||
return output_dict
|
||||
@@ -0,0 +1,179 @@
|
||||
from torch import nn
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from src.facerender.sync_batchnorm import SynchronizedBatchNorm2d as BatchNorm2d
|
||||
from src.facerender.modules.util import KPHourglass, make_coordinate_grid, AntiAliasInterpolation2d, ResBottleneck
|
||||
|
||||
|
||||
class KPDetector(nn.Module):
|
||||
"""
|
||||
Detecting canonical keypoints. Return keypoint position and jacobian near each keypoint.
|
||||
"""
|
||||
|
||||
def __init__(self, block_expansion, feature_channel, num_kp, image_channel, max_features, reshape_channel, reshape_depth,
|
||||
num_blocks, temperature, estimate_jacobian=False, scale_factor=1, single_jacobian_map=False):
|
||||
super(KPDetector, self).__init__()
|
||||
|
||||
self.predictor = KPHourglass(block_expansion, in_features=image_channel,
|
||||
max_features=max_features, reshape_features=reshape_channel, reshape_depth=reshape_depth, num_blocks=num_blocks)
|
||||
|
||||
# self.kp = nn.Conv3d(in_channels=self.predictor.out_filters, out_channels=num_kp, kernel_size=7, padding=3)
|
||||
self.kp = nn.Conv3d(in_channels=self.predictor.out_filters, out_channels=num_kp, kernel_size=3, padding=1)
|
||||
|
||||
if estimate_jacobian:
|
||||
self.num_jacobian_maps = 1 if single_jacobian_map else num_kp
|
||||
# self.jacobian = nn.Conv3d(in_channels=self.predictor.out_filters, out_channels=9 * self.num_jacobian_maps, kernel_size=7, padding=3)
|
||||
self.jacobian = nn.Conv3d(in_channels=self.predictor.out_filters, out_channels=9 * self.num_jacobian_maps, kernel_size=3, padding=1)
|
||||
'''
|
||||
initial as:
|
||||
[[1 0 0]
|
||||
[0 1 0]
|
||||
[0 0 1]]
|
||||
'''
|
||||
self.jacobian.weight.data.zero_()
|
||||
self.jacobian.bias.data.copy_(torch.tensor([1, 0, 0, 0, 1, 0, 0, 0, 1] * self.num_jacobian_maps, dtype=torch.float))
|
||||
else:
|
||||
self.jacobian = None
|
||||
|
||||
self.temperature = temperature
|
||||
self.scale_factor = scale_factor
|
||||
if self.scale_factor != 1:
|
||||
self.down = AntiAliasInterpolation2d(image_channel, self.scale_factor)
|
||||
|
||||
def gaussian2kp(self, heatmap):
|
||||
"""
|
||||
Extract the mean from a heatmap
|
||||
"""
|
||||
shape = heatmap.shape
|
||||
heatmap = heatmap.unsqueeze(-1)
|
||||
grid = make_coordinate_grid(shape[2:], heatmap.type()).unsqueeze_(0).unsqueeze_(0)
|
||||
value = (heatmap * grid).sum(dim=(2, 3, 4))
|
||||
kp = {'value': value}
|
||||
|
||||
return kp
|
||||
|
||||
def forward(self, x):
|
||||
if self.scale_factor != 1:
|
||||
x = self.down(x)
|
||||
|
||||
feature_map = self.predictor(x)
|
||||
prediction = self.kp(feature_map)
|
||||
|
||||
final_shape = prediction.shape
|
||||
heatmap = prediction.view(final_shape[0], final_shape[1], -1)
|
||||
heatmap = F.softmax(heatmap / self.temperature, dim=2)
|
||||
heatmap = heatmap.view(*final_shape)
|
||||
|
||||
out = self.gaussian2kp(heatmap)
|
||||
|
||||
if self.jacobian is not None:
|
||||
jacobian_map = self.jacobian(feature_map)
|
||||
jacobian_map = jacobian_map.reshape(final_shape[0], self.num_jacobian_maps, 9, final_shape[2],
|
||||
final_shape[3], final_shape[4])
|
||||
heatmap = heatmap.unsqueeze(2)
|
||||
|
||||
jacobian = heatmap * jacobian_map
|
||||
jacobian = jacobian.view(final_shape[0], final_shape[1], 9, -1)
|
||||
jacobian = jacobian.sum(dim=-1)
|
||||
jacobian = jacobian.view(jacobian.shape[0], jacobian.shape[1], 3, 3)
|
||||
out['jacobian'] = jacobian
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class HEEstimator(nn.Module):
|
||||
"""
|
||||
Estimating head pose and expression.
|
||||
"""
|
||||
|
||||
def __init__(self, block_expansion, feature_channel, num_kp, image_channel, max_features, num_bins=66, estimate_jacobian=True):
|
||||
super(HEEstimator, self).__init__()
|
||||
|
||||
self.conv1 = nn.Conv2d(in_channels=image_channel, out_channels=block_expansion, kernel_size=7, padding=3, stride=2)
|
||||
self.norm1 = BatchNorm2d(block_expansion, affine=True)
|
||||
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
|
||||
|
||||
self.conv2 = nn.Conv2d(in_channels=block_expansion, out_channels=256, kernel_size=1)
|
||||
self.norm2 = BatchNorm2d(256, affine=True)
|
||||
|
||||
self.block1 = nn.Sequential()
|
||||
for i in range(3):
|
||||
self.block1.add_module('b1_'+ str(i), ResBottleneck(in_features=256, stride=1))
|
||||
|
||||
self.conv3 = nn.Conv2d(in_channels=256, out_channels=512, kernel_size=1)
|
||||
self.norm3 = BatchNorm2d(512, affine=True)
|
||||
self.block2 = ResBottleneck(in_features=512, stride=2)
|
||||
|
||||
self.block3 = nn.Sequential()
|
||||
for i in range(3):
|
||||
self.block3.add_module('b3_'+ str(i), ResBottleneck(in_features=512, stride=1))
|
||||
|
||||
self.conv4 = nn.Conv2d(in_channels=512, out_channels=1024, kernel_size=1)
|
||||
self.norm4 = BatchNorm2d(1024, affine=True)
|
||||
self.block4 = ResBottleneck(in_features=1024, stride=2)
|
||||
|
||||
self.block5 = nn.Sequential()
|
||||
for i in range(5):
|
||||
self.block5.add_module('b5_'+ str(i), ResBottleneck(in_features=1024, stride=1))
|
||||
|
||||
self.conv5 = nn.Conv2d(in_channels=1024, out_channels=2048, kernel_size=1)
|
||||
self.norm5 = BatchNorm2d(2048, affine=True)
|
||||
self.block6 = ResBottleneck(in_features=2048, stride=2)
|
||||
|
||||
self.block7 = nn.Sequential()
|
||||
for i in range(2):
|
||||
self.block7.add_module('b7_'+ str(i), ResBottleneck(in_features=2048, stride=1))
|
||||
|
||||
self.fc_roll = nn.Linear(2048, num_bins)
|
||||
self.fc_pitch = nn.Linear(2048, num_bins)
|
||||
self.fc_yaw = nn.Linear(2048, num_bins)
|
||||
|
||||
self.fc_t = nn.Linear(2048, 3)
|
||||
|
||||
self.fc_exp = nn.Linear(2048, 3*num_kp)
|
||||
|
||||
def forward(self, x):
|
||||
out = self.conv1(x)
|
||||
out = self.norm1(out)
|
||||
out = F.relu(out)
|
||||
out = self.maxpool(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.norm2(out)
|
||||
out = F.relu(out)
|
||||
|
||||
out = self.block1(out)
|
||||
|
||||
out = self.conv3(out)
|
||||
out = self.norm3(out)
|
||||
out = F.relu(out)
|
||||
out = self.block2(out)
|
||||
|
||||
out = self.block3(out)
|
||||
|
||||
out = self.conv4(out)
|
||||
out = self.norm4(out)
|
||||
out = F.relu(out)
|
||||
out = self.block4(out)
|
||||
|
||||
out = self.block5(out)
|
||||
|
||||
out = self.conv5(out)
|
||||
out = self.norm5(out)
|
||||
out = F.relu(out)
|
||||
out = self.block6(out)
|
||||
|
||||
out = self.block7(out)
|
||||
|
||||
out = F.adaptive_avg_pool2d(out, 1)
|
||||
out = out.view(out.shape[0], -1)
|
||||
|
||||
yaw = self.fc_roll(out)
|
||||
pitch = self.fc_pitch(out)
|
||||
roll = self.fc_yaw(out)
|
||||
t = self.fc_t(out)
|
||||
exp = self.fc_exp(out)
|
||||
|
||||
return {'yaw': yaw, 'pitch': pitch, 'roll': roll, 't': t, 'exp': exp}
|
||||
|
||||
@@ -0,0 +1,170 @@
|
||||
from scipy.spatial import ConvexHull
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
from tqdm import tqdm
|
||||
|
||||
def normalize_kp(kp_source, kp_driving, kp_driving_initial, adapt_movement_scale=False,
|
||||
use_relative_movement=False, use_relative_jacobian=False):
|
||||
if adapt_movement_scale:
|
||||
source_area = ConvexHull(kp_source['value'][0].data.cpu().numpy()).volume
|
||||
driving_area = ConvexHull(kp_driving_initial['value'][0].data.cpu().numpy()).volume
|
||||
adapt_movement_scale = np.sqrt(source_area) / np.sqrt(driving_area)
|
||||
else:
|
||||
adapt_movement_scale = 1
|
||||
|
||||
kp_new = {k: v for k, v in kp_driving.items()}
|
||||
|
||||
if use_relative_movement:
|
||||
kp_value_diff = (kp_driving['value'] - kp_driving_initial['value'])
|
||||
kp_value_diff *= adapt_movement_scale
|
||||
kp_new['value'] = kp_value_diff + kp_source['value']
|
||||
|
||||
if use_relative_jacobian:
|
||||
jacobian_diff = torch.matmul(kp_driving['jacobian'], torch.inverse(kp_driving_initial['jacobian']))
|
||||
kp_new['jacobian'] = torch.matmul(jacobian_diff, kp_source['jacobian'])
|
||||
|
||||
return kp_new
|
||||
|
||||
def headpose_pred_to_degree(pred):
|
||||
device = pred.device
|
||||
idx_tensor = [idx for idx in range(66)]
|
||||
idx_tensor = torch.FloatTensor(idx_tensor).type_as(pred).to(device)
|
||||
pred = F.softmax(pred)
|
||||
degree = torch.sum(pred*idx_tensor, 1) * 3 - 99
|
||||
return degree
|
||||
|
||||
def get_rotation_matrix(yaw, pitch, roll):
|
||||
yaw = yaw / 180 * 3.14
|
||||
pitch = pitch / 180 * 3.14
|
||||
roll = roll / 180 * 3.14
|
||||
|
||||
roll = roll.unsqueeze(1)
|
||||
pitch = pitch.unsqueeze(1)
|
||||
yaw = yaw.unsqueeze(1)
|
||||
|
||||
pitch_mat = torch.cat([torch.ones_like(pitch), torch.zeros_like(pitch), torch.zeros_like(pitch),
|
||||
torch.zeros_like(pitch), torch.cos(pitch), -torch.sin(pitch),
|
||||
torch.zeros_like(pitch), torch.sin(pitch), torch.cos(pitch)], dim=1)
|
||||
pitch_mat = pitch_mat.view(pitch_mat.shape[0], 3, 3)
|
||||
|
||||
yaw_mat = torch.cat([torch.cos(yaw), torch.zeros_like(yaw), torch.sin(yaw),
|
||||
torch.zeros_like(yaw), torch.ones_like(yaw), torch.zeros_like(yaw),
|
||||
-torch.sin(yaw), torch.zeros_like(yaw), torch.cos(yaw)], dim=1)
|
||||
yaw_mat = yaw_mat.view(yaw_mat.shape[0], 3, 3)
|
||||
|
||||
roll_mat = torch.cat([torch.cos(roll), -torch.sin(roll), torch.zeros_like(roll),
|
||||
torch.sin(roll), torch.cos(roll), torch.zeros_like(roll),
|
||||
torch.zeros_like(roll), torch.zeros_like(roll), torch.ones_like(roll)], dim=1)
|
||||
roll_mat = roll_mat.view(roll_mat.shape[0], 3, 3)
|
||||
|
||||
rot_mat = torch.einsum('bij,bjk,bkm->bim', pitch_mat, yaw_mat, roll_mat)
|
||||
|
||||
return rot_mat
|
||||
|
||||
def keypoint_transformation(kp_canonical, he, wo_exp=False):
|
||||
kp = kp_canonical['value'] # (bs, k, 3)
|
||||
yaw, pitch, roll= he['yaw'], he['pitch'], he['roll']
|
||||
yaw = headpose_pred_to_degree(yaw)
|
||||
pitch = headpose_pred_to_degree(pitch)
|
||||
roll = headpose_pred_to_degree(roll)
|
||||
|
||||
if 'yaw_in' in he:
|
||||
yaw = he['yaw_in']
|
||||
if 'pitch_in' in he:
|
||||
pitch = he['pitch_in']
|
||||
if 'roll_in' in he:
|
||||
roll = he['roll_in']
|
||||
|
||||
rot_mat = get_rotation_matrix(yaw, pitch, roll) # (bs, 3, 3)
|
||||
|
||||
t, exp = he['t'], he['exp']
|
||||
if wo_exp:
|
||||
exp = exp*0
|
||||
|
||||
# keypoint rotation
|
||||
kp_rotated = torch.einsum('bmp,bkp->bkm', rot_mat, kp)
|
||||
|
||||
# keypoint translation
|
||||
t[:, 0] = t[:, 0]*0
|
||||
t[:, 2] = t[:, 2]*0
|
||||
t = t.unsqueeze(1).repeat(1, kp.shape[1], 1)
|
||||
kp_t = kp_rotated + t
|
||||
|
||||
# add expression deviation
|
||||
exp = exp.view(exp.shape[0], -1, 3)
|
||||
kp_transformed = kp_t + exp
|
||||
|
||||
return {'value': kp_transformed}
|
||||
|
||||
|
||||
|
||||
def make_animation(source_image, source_semantics, target_semantics,
|
||||
generator, kp_detector, he_estimator, mapping,
|
||||
yaw_c_seq=None, pitch_c_seq=None, roll_c_seq=None,
|
||||
use_exp=True, use_half=False):
|
||||
with torch.no_grad():
|
||||
predictions = []
|
||||
|
||||
kp_canonical = kp_detector(source_image)
|
||||
he_source = mapping(source_semantics)
|
||||
kp_source = keypoint_transformation(kp_canonical, he_source)
|
||||
|
||||
for frame_idx in tqdm(range(target_semantics.shape[1]), 'Face Renderer:'):
|
||||
# still check the dimension
|
||||
# print(target_semantics.shape, source_semantics.shape)
|
||||
target_semantics_frame = target_semantics[:, frame_idx]
|
||||
he_driving = mapping(target_semantics_frame)
|
||||
if yaw_c_seq is not None:
|
||||
he_driving['yaw_in'] = yaw_c_seq[:, frame_idx]
|
||||
if pitch_c_seq is not None:
|
||||
he_driving['pitch_in'] = pitch_c_seq[:, frame_idx]
|
||||
if roll_c_seq is not None:
|
||||
he_driving['roll_in'] = roll_c_seq[:, frame_idx]
|
||||
|
||||
kp_driving = keypoint_transformation(kp_canonical, he_driving)
|
||||
|
||||
kp_norm = kp_driving
|
||||
out = generator(source_image, kp_source=kp_source, kp_driving=kp_norm)
|
||||
'''
|
||||
source_image_new = out['prediction'].squeeze(1)
|
||||
kp_canonical_new = kp_detector(source_image_new)
|
||||
he_source_new = he_estimator(source_image_new)
|
||||
kp_source_new = keypoint_transformation(kp_canonical_new, he_source_new, wo_exp=True)
|
||||
kp_driving_new = keypoint_transformation(kp_canonical_new, he_driving, wo_exp=True)
|
||||
out = generator(source_image_new, kp_source=kp_source_new, kp_driving=kp_driving_new)
|
||||
'''
|
||||
predictions.append(out['prediction'])
|
||||
predictions_ts = torch.stack(predictions, dim=1)
|
||||
return predictions_ts
|
||||
|
||||
class AnimateModel(torch.nn.Module):
|
||||
"""
|
||||
Merge all generator related updates into single model for better multi-gpu usage
|
||||
"""
|
||||
|
||||
def __init__(self, generator, kp_extractor, mapping):
|
||||
super(AnimateModel, self).__init__()
|
||||
self.kp_extractor = kp_extractor
|
||||
self.generator = generator
|
||||
self.mapping = mapping
|
||||
|
||||
self.kp_extractor.eval()
|
||||
self.generator.eval()
|
||||
self.mapping.eval()
|
||||
|
||||
def forward(self, x):
|
||||
|
||||
source_image = x['source_image']
|
||||
source_semantics = x['source_semantics']
|
||||
target_semantics = x['target_semantics']
|
||||
yaw_c_seq = x['yaw_c_seq']
|
||||
pitch_c_seq = x['pitch_c_seq']
|
||||
roll_c_seq = x['roll_c_seq']
|
||||
|
||||
predictions_video = make_animation(source_image, source_semantics, target_semantics,
|
||||
self.generator, self.kp_extractor,
|
||||
self.mapping, use_exp = True,
|
||||
yaw_c_seq=yaw_c_seq, pitch_c_seq=pitch_c_seq, roll_c_seq=roll_c_seq)
|
||||
|
||||
return predictions_video
|
||||
@@ -0,0 +1,47 @@
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class MappingNet(nn.Module):
|
||||
def __init__(self, coeff_nc, descriptor_nc, layer, num_kp, num_bins):
|
||||
super( MappingNet, self).__init__()
|
||||
|
||||
self.layer = layer
|
||||
nonlinearity = nn.LeakyReLU(0.1)
|
||||
|
||||
self.first = nn.Sequential(
|
||||
torch.nn.Conv1d(coeff_nc, descriptor_nc, kernel_size=7, padding=0, bias=True))
|
||||
|
||||
for i in range(layer):
|
||||
net = nn.Sequential(nonlinearity,
|
||||
torch.nn.Conv1d(descriptor_nc, descriptor_nc, kernel_size=3, padding=0, dilation=3))
|
||||
setattr(self, 'encoder' + str(i), net)
|
||||
|
||||
self.pooling = nn.AdaptiveAvgPool1d(1)
|
||||
self.output_nc = descriptor_nc
|
||||
|
||||
self.fc_roll = nn.Linear(descriptor_nc, num_bins)
|
||||
self.fc_pitch = nn.Linear(descriptor_nc, num_bins)
|
||||
self.fc_yaw = nn.Linear(descriptor_nc, num_bins)
|
||||
self.fc_t = nn.Linear(descriptor_nc, 3)
|
||||
self.fc_exp = nn.Linear(descriptor_nc, 3*num_kp)
|
||||
|
||||
def forward(self, input_3dmm):
|
||||
out = self.first(input_3dmm)
|
||||
for i in range(self.layer):
|
||||
model = getattr(self, 'encoder' + str(i))
|
||||
out = model(out) + out[:,:,3:-3]
|
||||
out = self.pooling(out)
|
||||
out = out.view(out.shape[0], -1)
|
||||
#print('out:', out.shape)
|
||||
|
||||
yaw = self.fc_yaw(out)
|
||||
pitch = self.fc_pitch(out)
|
||||
roll = self.fc_roll(out)
|
||||
t = self.fc_t(out)
|
||||
exp = self.fc_exp(out)
|
||||
|
||||
return {'yaw': yaw, 'pitch': pitch, 'roll': roll, 't': t, 'exp': exp}
|
||||
@@ -0,0 +1,564 @@
|
||||
from torch import nn
|
||||
|
||||
import torch.nn.functional as F
|
||||
import torch
|
||||
|
||||
from src.facerender.sync_batchnorm import SynchronizedBatchNorm2d as BatchNorm2d
|
||||
from src.facerender.sync_batchnorm import SynchronizedBatchNorm3d as BatchNorm3d
|
||||
|
||||
import torch.nn.utils.spectral_norm as spectral_norm
|
||||
|
||||
|
||||
def kp2gaussian(kp, spatial_size, kp_variance):
|
||||
"""
|
||||
Transform a keypoint into gaussian like representation
|
||||
"""
|
||||
mean = kp['value']
|
||||
|
||||
coordinate_grid = make_coordinate_grid(spatial_size, mean.type())
|
||||
number_of_leading_dimensions = len(mean.shape) - 1
|
||||
shape = (1,) * number_of_leading_dimensions + coordinate_grid.shape
|
||||
coordinate_grid = coordinate_grid.view(*shape)
|
||||
repeats = mean.shape[:number_of_leading_dimensions] + (1, 1, 1, 1)
|
||||
coordinate_grid = coordinate_grid.repeat(*repeats)
|
||||
|
||||
# Preprocess kp shape
|
||||
shape = mean.shape[:number_of_leading_dimensions] + (1, 1, 1, 3)
|
||||
mean = mean.view(*shape)
|
||||
|
||||
mean_sub = (coordinate_grid - mean)
|
||||
|
||||
out = torch.exp(-0.5 * (mean_sub ** 2).sum(-1) / kp_variance)
|
||||
|
||||
return out
|
||||
|
||||
def make_coordinate_grid_2d(spatial_size, type):
|
||||
"""
|
||||
Create a meshgrid [-1,1] x [-1,1] of given spatial_size.
|
||||
"""
|
||||
h, w = spatial_size
|
||||
x = torch.arange(w).type(type)
|
||||
y = torch.arange(h).type(type)
|
||||
|
||||
x = (2 * (x / (w - 1)) - 1)
|
||||
y = (2 * (y / (h - 1)) - 1)
|
||||
|
||||
yy = y.view(-1, 1).repeat(1, w)
|
||||
xx = x.view(1, -1).repeat(h, 1)
|
||||
|
||||
meshed = torch.cat([xx.unsqueeze_(2), yy.unsqueeze_(2)], 2)
|
||||
|
||||
return meshed
|
||||
|
||||
|
||||
def make_coordinate_grid(spatial_size, type):
|
||||
d, h, w = spatial_size
|
||||
x = torch.arange(w).type(type)
|
||||
y = torch.arange(h).type(type)
|
||||
z = torch.arange(d).type(type)
|
||||
|
||||
x = (2 * (x / (w - 1)) - 1)
|
||||
y = (2 * (y / (h - 1)) - 1)
|
||||
z = (2 * (z / (d - 1)) - 1)
|
||||
|
||||
yy = y.view(1, -1, 1).repeat(d, 1, w)
|
||||
xx = x.view(1, 1, -1).repeat(d, h, 1)
|
||||
zz = z.view(-1, 1, 1).repeat(1, h, w)
|
||||
|
||||
meshed = torch.cat([xx.unsqueeze_(3), yy.unsqueeze_(3), zz.unsqueeze_(3)], 3)
|
||||
|
||||
return meshed
|
||||
|
||||
|
||||
class ResBottleneck(nn.Module):
|
||||
def __init__(self, in_features, stride):
|
||||
super(ResBottleneck, self).__init__()
|
||||
self.conv1 = nn.Conv2d(in_channels=in_features, out_channels=in_features//4, kernel_size=1)
|
||||
self.conv2 = nn.Conv2d(in_channels=in_features//4, out_channels=in_features//4, kernel_size=3, padding=1, stride=stride)
|
||||
self.conv3 = nn.Conv2d(in_channels=in_features//4, out_channels=in_features, kernel_size=1)
|
||||
self.norm1 = BatchNorm2d(in_features//4, affine=True)
|
||||
self.norm2 = BatchNorm2d(in_features//4, affine=True)
|
||||
self.norm3 = BatchNorm2d(in_features, affine=True)
|
||||
|
||||
self.stride = stride
|
||||
if self.stride != 1:
|
||||
self.skip = nn.Conv2d(in_channels=in_features, out_channels=in_features, kernel_size=1, stride=stride)
|
||||
self.norm4 = BatchNorm2d(in_features, affine=True)
|
||||
|
||||
def forward(self, x):
|
||||
out = self.conv1(x)
|
||||
out = self.norm1(out)
|
||||
out = F.relu(out)
|
||||
out = self.conv2(out)
|
||||
out = self.norm2(out)
|
||||
out = F.relu(out)
|
||||
out = self.conv3(out)
|
||||
out = self.norm3(out)
|
||||
if self.stride != 1:
|
||||
x = self.skip(x)
|
||||
x = self.norm4(x)
|
||||
out += x
|
||||
out = F.relu(out)
|
||||
return out
|
||||
|
||||
|
||||
class ResBlock2d(nn.Module):
|
||||
"""
|
||||
Res block, preserve spatial resolution.
|
||||
"""
|
||||
|
||||
def __init__(self, in_features, kernel_size, padding):
|
||||
super(ResBlock2d, self).__init__()
|
||||
self.conv1 = nn.Conv2d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size,
|
||||
padding=padding)
|
||||
self.conv2 = nn.Conv2d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size,
|
||||
padding=padding)
|
||||
self.norm1 = BatchNorm2d(in_features, affine=True)
|
||||
self.norm2 = BatchNorm2d(in_features, affine=True)
|
||||
|
||||
def forward(self, x):
|
||||
out = self.norm1(x)
|
||||
out = F.relu(out)
|
||||
out = self.conv1(out)
|
||||
out = self.norm2(out)
|
||||
out = F.relu(out)
|
||||
out = self.conv2(out)
|
||||
out += x
|
||||
return out
|
||||
|
||||
|
||||
class ResBlock3d(nn.Module):
|
||||
"""
|
||||
Res block, preserve spatial resolution.
|
||||
"""
|
||||
|
||||
def __init__(self, in_features, kernel_size, padding):
|
||||
super(ResBlock3d, self).__init__()
|
||||
self.conv1 = nn.Conv3d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size,
|
||||
padding=padding)
|
||||
self.conv2 = nn.Conv3d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size,
|
||||
padding=padding)
|
||||
self.norm1 = BatchNorm3d(in_features, affine=True)
|
||||
self.norm2 = BatchNorm3d(in_features, affine=True)
|
||||
|
||||
def forward(self, x):
|
||||
out = self.norm1(x)
|
||||
out = F.relu(out)
|
||||
out = self.conv1(out)
|
||||
out = self.norm2(out)
|
||||
out = F.relu(out)
|
||||
out = self.conv2(out)
|
||||
out += x
|
||||
return out
|
||||
|
||||
|
||||
class UpBlock2d(nn.Module):
|
||||
"""
|
||||
Upsampling block for use in decoder.
|
||||
"""
|
||||
|
||||
def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1):
|
||||
super(UpBlock2d, self).__init__()
|
||||
|
||||
self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size,
|
||||
padding=padding, groups=groups)
|
||||
self.norm = BatchNorm2d(out_features, affine=True)
|
||||
|
||||
def forward(self, x):
|
||||
out = F.interpolate(x, scale_factor=2)
|
||||
out = self.conv(out)
|
||||
out = self.norm(out)
|
||||
out = F.relu(out)
|
||||
return out
|
||||
|
||||
class UpBlock3d(nn.Module):
|
||||
"""
|
||||
Upsampling block for use in decoder.
|
||||
"""
|
||||
|
||||
def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1):
|
||||
super(UpBlock3d, self).__init__()
|
||||
|
||||
self.conv = nn.Conv3d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size,
|
||||
padding=padding, groups=groups)
|
||||
self.norm = BatchNorm3d(out_features, affine=True)
|
||||
|
||||
def forward(self, x):
|
||||
# out = F.interpolate(x, scale_factor=(1, 2, 2), mode='trilinear')
|
||||
out = F.interpolate(x, scale_factor=(1, 2, 2))
|
||||
out = self.conv(out)
|
||||
out = self.norm(out)
|
||||
out = F.relu(out)
|
||||
return out
|
||||
|
||||
|
||||
class DownBlock2d(nn.Module):
|
||||
"""
|
||||
Downsampling block for use in encoder.
|
||||
"""
|
||||
|
||||
def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1):
|
||||
super(DownBlock2d, self).__init__()
|
||||
self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size,
|
||||
padding=padding, groups=groups)
|
||||
self.norm = BatchNorm2d(out_features, affine=True)
|
||||
self.pool = nn.AvgPool2d(kernel_size=(2, 2))
|
||||
|
||||
def forward(self, x):
|
||||
out = self.conv(x)
|
||||
out = self.norm(out)
|
||||
out = F.relu(out)
|
||||
out = self.pool(out)
|
||||
return out
|
||||
|
||||
|
||||
class DownBlock3d(nn.Module):
|
||||
"""
|
||||
Downsampling block for use in encoder.
|
||||
"""
|
||||
|
||||
def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1):
|
||||
super(DownBlock3d, self).__init__()
|
||||
'''
|
||||
self.conv = nn.Conv3d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size,
|
||||
padding=padding, groups=groups, stride=(1, 2, 2))
|
||||
'''
|
||||
self.conv = nn.Conv3d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size,
|
||||
padding=padding, groups=groups)
|
||||
self.norm = BatchNorm3d(out_features, affine=True)
|
||||
self.pool = nn.AvgPool3d(kernel_size=(1, 2, 2))
|
||||
|
||||
def forward(self, x):
|
||||
out = self.conv(x)
|
||||
out = self.norm(out)
|
||||
out = F.relu(out)
|
||||
out = self.pool(out)
|
||||
return out
|
||||
|
||||
|
||||
class SameBlock2d(nn.Module):
|
||||
"""
|
||||
Simple block, preserve spatial resolution.
|
||||
"""
|
||||
|
||||
def __init__(self, in_features, out_features, groups=1, kernel_size=3, padding=1, lrelu=False):
|
||||
super(SameBlock2d, self).__init__()
|
||||
self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features,
|
||||
kernel_size=kernel_size, padding=padding, groups=groups)
|
||||
self.norm = BatchNorm2d(out_features, affine=True)
|
||||
if lrelu:
|
||||
self.ac = nn.LeakyReLU()
|
||||
else:
|
||||
self.ac = nn.ReLU()
|
||||
|
||||
def forward(self, x):
|
||||
out = self.conv(x)
|
||||
out = self.norm(out)
|
||||
out = self.ac(out)
|
||||
return out
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
"""
|
||||
Hourglass Encoder
|
||||
"""
|
||||
|
||||
def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256):
|
||||
super(Encoder, self).__init__()
|
||||
|
||||
down_blocks = []
|
||||
for i in range(num_blocks):
|
||||
down_blocks.append(DownBlock3d(in_features if i == 0 else min(max_features, block_expansion * (2 ** i)),
|
||||
min(max_features, block_expansion * (2 ** (i + 1))),
|
||||
kernel_size=3, padding=1))
|
||||
self.down_blocks = nn.ModuleList(down_blocks)
|
||||
|
||||
def forward(self, x):
|
||||
outs = [x]
|
||||
for down_block in self.down_blocks:
|
||||
outs.append(down_block(outs[-1]))
|
||||
return outs
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
"""
|
||||
Hourglass Decoder
|
||||
"""
|
||||
|
||||
def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256):
|
||||
super(Decoder, self).__init__()
|
||||
|
||||
up_blocks = []
|
||||
|
||||
for i in range(num_blocks)[::-1]:
|
||||
in_filters = (1 if i == num_blocks - 1 else 2) * min(max_features, block_expansion * (2 ** (i + 1)))
|
||||
out_filters = min(max_features, block_expansion * (2 ** i))
|
||||
up_blocks.append(UpBlock3d(in_filters, out_filters, kernel_size=3, padding=1))
|
||||
|
||||
self.up_blocks = nn.ModuleList(up_blocks)
|
||||
# self.out_filters = block_expansion
|
||||
self.out_filters = block_expansion + in_features
|
||||
|
||||
self.conv = nn.Conv3d(in_channels=self.out_filters, out_channels=self.out_filters, kernel_size=3, padding=1)
|
||||
self.norm = BatchNorm3d(self.out_filters, affine=True)
|
||||
|
||||
def forward(self, x):
|
||||
out = x.pop()
|
||||
# for up_block in self.up_blocks[:-1]:
|
||||
for up_block in self.up_blocks:
|
||||
out = up_block(out)
|
||||
skip = x.pop()
|
||||
out = torch.cat([out, skip], dim=1)
|
||||
# out = self.up_blocks[-1](out)
|
||||
out = self.conv(out)
|
||||
out = self.norm(out)
|
||||
out = F.relu(out)
|
||||
return out
|
||||
|
||||
|
||||
class Hourglass(nn.Module):
|
||||
"""
|
||||
Hourglass architecture.
|
||||
"""
|
||||
|
||||
def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256):
|
||||
super(Hourglass, self).__init__()
|
||||
self.encoder = Encoder(block_expansion, in_features, num_blocks, max_features)
|
||||
self.decoder = Decoder(block_expansion, in_features, num_blocks, max_features)
|
||||
self.out_filters = self.decoder.out_filters
|
||||
|
||||
def forward(self, x):
|
||||
return self.decoder(self.encoder(x))
|
||||
|
||||
|
||||
class KPHourglass(nn.Module):
|
||||
"""
|
||||
Hourglass architecture.
|
||||
"""
|
||||
|
||||
def __init__(self, block_expansion, in_features, reshape_features, reshape_depth, num_blocks=3, max_features=256):
|
||||
super(KPHourglass, self).__init__()
|
||||
|
||||
self.down_blocks = nn.Sequential()
|
||||
for i in range(num_blocks):
|
||||
self.down_blocks.add_module('down'+ str(i), DownBlock2d(in_features if i == 0 else min(max_features, block_expansion * (2 ** i)),
|
||||
min(max_features, block_expansion * (2 ** (i + 1))),
|
||||
kernel_size=3, padding=1))
|
||||
|
||||
in_filters = min(max_features, block_expansion * (2 ** num_blocks))
|
||||
self.conv = nn.Conv2d(in_channels=in_filters, out_channels=reshape_features, kernel_size=1)
|
||||
|
||||
self.up_blocks = nn.Sequential()
|
||||
for i in range(num_blocks):
|
||||
in_filters = min(max_features, block_expansion * (2 ** (num_blocks - i)))
|
||||
out_filters = min(max_features, block_expansion * (2 ** (num_blocks - i - 1)))
|
||||
self.up_blocks.add_module('up'+ str(i), UpBlock3d(in_filters, out_filters, kernel_size=3, padding=1))
|
||||
|
||||
self.reshape_depth = reshape_depth
|
||||
self.out_filters = out_filters
|
||||
|
||||
def forward(self, x):
|
||||
out = self.down_blocks(x)
|
||||
out = self.conv(out)
|
||||
bs, c, h, w = out.shape
|
||||
out = out.view(bs, c//self.reshape_depth, self.reshape_depth, h, w)
|
||||
out = self.up_blocks(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
|
||||
class AntiAliasInterpolation2d(nn.Module):
|
||||
"""
|
||||
Band-limited downsampling, for better preservation of the input signal.
|
||||
"""
|
||||
def __init__(self, channels, scale):
|
||||
super(AntiAliasInterpolation2d, self).__init__()
|
||||
sigma = (1 / scale - 1) / 2
|
||||
kernel_size = 2 * round(sigma * 4) + 1
|
||||
self.ka = kernel_size // 2
|
||||
self.kb = self.ka - 1 if kernel_size % 2 == 0 else self.ka
|
||||
|
||||
kernel_size = [kernel_size, kernel_size]
|
||||
sigma = [sigma, sigma]
|
||||
# The gaussian kernel is the product of the
|
||||
# gaussian function of each dimension.
|
||||
kernel = 1
|
||||
meshgrids = torch.meshgrid(
|
||||
[
|
||||
torch.arange(size, dtype=torch.float32)
|
||||
for size in kernel_size
|
||||
]
|
||||
)
|
||||
for size, std, mgrid in zip(kernel_size, sigma, meshgrids):
|
||||
mean = (size - 1) / 2
|
||||
kernel *= torch.exp(-(mgrid - mean) ** 2 / (2 * std ** 2))
|
||||
|
||||
# Make sure sum of values in gaussian kernel equals 1.
|
||||
kernel = kernel / torch.sum(kernel)
|
||||
# Reshape to depthwise convolutional weight
|
||||
kernel = kernel.view(1, 1, *kernel.size())
|
||||
kernel = kernel.repeat(channels, *[1] * (kernel.dim() - 1))
|
||||
|
||||
self.register_buffer('weight', kernel)
|
||||
self.groups = channels
|
||||
self.scale = scale
|
||||
inv_scale = 1 / scale
|
||||
self.int_inv_scale = int(inv_scale)
|
||||
|
||||
def forward(self, input):
|
||||
if self.scale == 1.0:
|
||||
return input
|
||||
|
||||
out = F.pad(input, (self.ka, self.kb, self.ka, self.kb))
|
||||
out = F.conv2d(out, weight=self.weight, groups=self.groups)
|
||||
out = out[:, :, ::self.int_inv_scale, ::self.int_inv_scale]
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class SPADE(nn.Module):
|
||||
def __init__(self, norm_nc, label_nc):
|
||||
super().__init__()
|
||||
|
||||
self.param_free_norm = nn.InstanceNorm2d(norm_nc, affine=False)
|
||||
nhidden = 128
|
||||
|
||||
self.mlp_shared = nn.Sequential(
|
||||
nn.Conv2d(label_nc, nhidden, kernel_size=3, padding=1),
|
||||
nn.ReLU())
|
||||
self.mlp_gamma = nn.Conv2d(nhidden, norm_nc, kernel_size=3, padding=1)
|
||||
self.mlp_beta = nn.Conv2d(nhidden, norm_nc, kernel_size=3, padding=1)
|
||||
|
||||
def forward(self, x, segmap):
|
||||
normalized = self.param_free_norm(x)
|
||||
segmap = F.interpolate(segmap, size=x.size()[2:], mode='nearest')
|
||||
actv = self.mlp_shared(segmap)
|
||||
gamma = self.mlp_gamma(actv)
|
||||
beta = self.mlp_beta(actv)
|
||||
out = normalized * (1 + gamma) + beta
|
||||
return out
|
||||
|
||||
|
||||
class SPADEResnetBlock(nn.Module):
|
||||
def __init__(self, fin, fout, norm_G, label_nc, use_se=False, dilation=1):
|
||||
super().__init__()
|
||||
# Attributes
|
||||
self.learned_shortcut = (fin != fout)
|
||||
fmiddle = min(fin, fout)
|
||||
self.use_se = use_se
|
||||
# create conv layers
|
||||
self.conv_0 = nn.Conv2d(fin, fmiddle, kernel_size=3, padding=dilation, dilation=dilation)
|
||||
self.conv_1 = nn.Conv2d(fmiddle, fout, kernel_size=3, padding=dilation, dilation=dilation)
|
||||
if self.learned_shortcut:
|
||||
self.conv_s = nn.Conv2d(fin, fout, kernel_size=1, bias=False)
|
||||
# apply spectral norm if specified
|
||||
if 'spectral' in norm_G:
|
||||
self.conv_0 = spectral_norm(self.conv_0)
|
||||
self.conv_1 = spectral_norm(self.conv_1)
|
||||
if self.learned_shortcut:
|
||||
self.conv_s = spectral_norm(self.conv_s)
|
||||
# define normalization layers
|
||||
self.norm_0 = SPADE(fin, label_nc)
|
||||
self.norm_1 = SPADE(fmiddle, label_nc)
|
||||
if self.learned_shortcut:
|
||||
self.norm_s = SPADE(fin, label_nc)
|
||||
|
||||
def forward(self, x, seg1):
|
||||
x_s = self.shortcut(x, seg1)
|
||||
dx = self.conv_0(self.actvn(self.norm_0(x, seg1)))
|
||||
dx = self.conv_1(self.actvn(self.norm_1(dx, seg1)))
|
||||
out = x_s + dx
|
||||
return out
|
||||
|
||||
def shortcut(self, x, seg1):
|
||||
if self.learned_shortcut:
|
||||
x_s = self.conv_s(self.norm_s(x, seg1))
|
||||
else:
|
||||
x_s = x
|
||||
return x_s
|
||||
|
||||
def actvn(self, x):
|
||||
return F.leaky_relu(x, 2e-1)
|
||||
|
||||
class audio2image(nn.Module):
|
||||
def __init__(self, generator, kp_extractor, he_estimator_video, he_estimator_audio, train_params):
|
||||
super().__init__()
|
||||
# Attributes
|
||||
self.generator = generator
|
||||
self.kp_extractor = kp_extractor
|
||||
self.he_estimator_video = he_estimator_video
|
||||
self.he_estimator_audio = he_estimator_audio
|
||||
self.train_params = train_params
|
||||
|
||||
def headpose_pred_to_degree(self, pred):
|
||||
device = pred.device
|
||||
idx_tensor = [idx for idx in range(66)]
|
||||
idx_tensor = torch.FloatTensor(idx_tensor).to(device)
|
||||
pred = F.softmax(pred)
|
||||
degree = torch.sum(pred*idx_tensor, 1) * 3 - 99
|
||||
|
||||
return degree
|
||||
|
||||
def get_rotation_matrix(self, yaw, pitch, roll):
|
||||
yaw = yaw / 180 * 3.14
|
||||
pitch = pitch / 180 * 3.14
|
||||
roll = roll / 180 * 3.14
|
||||
|
||||
roll = roll.unsqueeze(1)
|
||||
pitch = pitch.unsqueeze(1)
|
||||
yaw = yaw.unsqueeze(1)
|
||||
|
||||
roll_mat = torch.cat([torch.ones_like(roll), torch.zeros_like(roll), torch.zeros_like(roll),
|
||||
torch.zeros_like(roll), torch.cos(roll), -torch.sin(roll),
|
||||
torch.zeros_like(roll), torch.sin(roll), torch.cos(roll)], dim=1)
|
||||
roll_mat = roll_mat.view(roll_mat.shape[0], 3, 3)
|
||||
|
||||
pitch_mat = torch.cat([torch.cos(pitch), torch.zeros_like(pitch), torch.sin(pitch),
|
||||
torch.zeros_like(pitch), torch.ones_like(pitch), torch.zeros_like(pitch),
|
||||
-torch.sin(pitch), torch.zeros_like(pitch), torch.cos(pitch)], dim=1)
|
||||
pitch_mat = pitch_mat.view(pitch_mat.shape[0], 3, 3)
|
||||
|
||||
yaw_mat = torch.cat([torch.cos(yaw), -torch.sin(yaw), torch.zeros_like(yaw),
|
||||
torch.sin(yaw), torch.cos(yaw), torch.zeros_like(yaw),
|
||||
torch.zeros_like(yaw), torch.zeros_like(yaw), torch.ones_like(yaw)], dim=1)
|
||||
yaw_mat = yaw_mat.view(yaw_mat.shape[0], 3, 3)
|
||||
|
||||
rot_mat = torch.einsum('bij,bjk,bkm->bim', roll_mat, pitch_mat, yaw_mat)
|
||||
|
||||
return rot_mat
|
||||
|
||||
def keypoint_transformation(self, kp_canonical, he):
|
||||
kp = kp_canonical['value'] # (bs, k, 3)
|
||||
yaw, pitch, roll = he['yaw'], he['pitch'], he['roll']
|
||||
t, exp = he['t'], he['exp']
|
||||
|
||||
yaw = self.headpose_pred_to_degree(yaw)
|
||||
pitch = self.headpose_pred_to_degree(pitch)
|
||||
roll = self.headpose_pred_to_degree(roll)
|
||||
|
||||
rot_mat = self.get_rotation_matrix(yaw, pitch, roll) # (bs, 3, 3)
|
||||
|
||||
# keypoint rotation
|
||||
kp_rotated = torch.einsum('bmp,bkp->bkm', rot_mat, kp)
|
||||
|
||||
|
||||
|
||||
# keypoint translation
|
||||
t = t.unsqueeze_(1).repeat(1, kp.shape[1], 1)
|
||||
kp_t = kp_rotated + t
|
||||
|
||||
# add expression deviation
|
||||
exp = exp.view(exp.shape[0], -1, 3)
|
||||
kp_transformed = kp_t + exp
|
||||
|
||||
return {'value': kp_transformed}
|
||||
|
||||
def forward(self, source_image, target_audio):
|
||||
pose_source = self.he_estimator_video(source_image)
|
||||
pose_generated = self.he_estimator_audio(target_audio)
|
||||
kp_canonical = self.kp_extractor(source_image)
|
||||
kp_source = self.keypoint_transformation(kp_canonical, pose_source)
|
||||
kp_transformed_generated = self.keypoint_transformation(kp_canonical, pose_generated)
|
||||
generated = self.generator(source_image, kp_source=kp_source, kp_driving=kp_transformed_generated)
|
||||
return generated
|
||||
Reference in New Issue
Block a user