Add files via upload
This commit is contained in:
VangengLab
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,48 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
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) # Bx3x256x256 -> Bx64x256x256
|
||||
|
||||
for i in range(len(self.down_blocks)):
|
||||
out = self.down_blocks[i](out)
|
||||
out = self.second(out)
|
||||
bs, c, h, w = out.shape # ->Bx512x64x64
|
||||
|
||||
f_s = out.view(bs, self.reshape_channel, self.reshape_depth, h, w) # ->Bx32x16x64x64
|
||||
f_s = self.resblocks_3d(f_s) # ->Bx32x16x64x64
|
||||
return f_s
|
||||
@@ -0,0 +1,149 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
This moudle is adapted to the ConvNeXtV2 version for the extraction of implicit keypoints, poses, and expression deformation.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
# from timm.models.layers import trunc_normal_, DropPath
|
||||
from .util import LayerNorm, DropPath, trunc_normal_, GRN
|
||||
|
||||
__all__ = ['convnextv2_tiny']
|
||||
|
||||
|
||||
class Block(nn.Module):
|
||||
""" ConvNeXtV2 Block.
|
||||
|
||||
Args:
|
||||
dim (int): Number of input channels.
|
||||
drop_path (float): Stochastic depth rate. Default: 0.0
|
||||
"""
|
||||
|
||||
def __init__(self, dim, drop_path=0.):
|
||||
super().__init__()
|
||||
self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) # depthwise conv
|
||||
self.norm = LayerNorm(dim, eps=1e-6)
|
||||
self.pwconv1 = nn.Linear(dim, 4 * dim) # pointwise/1x1 convs, implemented with linear layers
|
||||
self.act = nn.GELU()
|
||||
self.grn = GRN(4 * dim)
|
||||
self.pwconv2 = nn.Linear(4 * dim, dim)
|
||||
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
|
||||
|
||||
def forward(self, x):
|
||||
input = x
|
||||
x = self.dwconv(x)
|
||||
x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
|
||||
x = self.norm(x)
|
||||
x = self.pwconv1(x)
|
||||
x = self.act(x)
|
||||
x = self.grn(x)
|
||||
x = self.pwconv2(x)
|
||||
x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W)
|
||||
|
||||
x = input + self.drop_path(x)
|
||||
return x
|
||||
|
||||
|
||||
class ConvNeXtV2(nn.Module):
|
||||
""" ConvNeXt V2
|
||||
|
||||
Args:
|
||||
in_chans (int): Number of input image channels. Default: 3
|
||||
num_classes (int): Number of classes for classification head. Default: 1000
|
||||
depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
|
||||
dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
|
||||
drop_path_rate (float): Stochastic depth rate. Default: 0.
|
||||
head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_chans=3,
|
||||
depths=[3, 3, 9, 3],
|
||||
dims=[96, 192, 384, 768],
|
||||
drop_path_rate=0.,
|
||||
**kwargs
|
||||
):
|
||||
super().__init__()
|
||||
self.depths = depths
|
||||
self.downsample_layers = nn.ModuleList() # stem and 3 intermediate downsampling conv layers
|
||||
stem = nn.Sequential(
|
||||
nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4),
|
||||
LayerNorm(dims[0], eps=1e-6, data_format="channels_first")
|
||||
)
|
||||
self.downsample_layers.append(stem)
|
||||
for i in range(3):
|
||||
downsample_layer = nn.Sequential(
|
||||
LayerNorm(dims[i], eps=1e-6, data_format="channels_first"),
|
||||
nn.Conv2d(dims[i], dims[i+1], kernel_size=2, stride=2),
|
||||
)
|
||||
self.downsample_layers.append(downsample_layer)
|
||||
|
||||
self.stages = nn.ModuleList() # 4 feature resolution stages, each consisting of multiple residual blocks
|
||||
dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
|
||||
cur = 0
|
||||
for i in range(4):
|
||||
stage = nn.Sequential(
|
||||
*[Block(dim=dims[i], drop_path=dp_rates[cur + j]) for j in range(depths[i])]
|
||||
)
|
||||
self.stages.append(stage)
|
||||
cur += depths[i]
|
||||
|
||||
self.norm = nn.LayerNorm(dims[-1], eps=1e-6) # final norm layer
|
||||
|
||||
# NOTE: the output semantic items
|
||||
num_bins = kwargs.get('num_bins', 66)
|
||||
num_kp = kwargs.get('num_kp', 24) # the number of implicit keypoints
|
||||
self.fc_kp = nn.Linear(dims[-1], 3 * num_kp) # implicit keypoints
|
||||
|
||||
# print('dims[-1]: ', dims[-1])
|
||||
self.fc_scale = nn.Linear(dims[-1], 1) # scale
|
||||
self.fc_pitch = nn.Linear(dims[-1], num_bins) # pitch bins
|
||||
self.fc_yaw = nn.Linear(dims[-1], num_bins) # yaw bins
|
||||
self.fc_roll = nn.Linear(dims[-1], num_bins) # roll bins
|
||||
self.fc_t = nn.Linear(dims[-1], 3) # translation
|
||||
self.fc_exp = nn.Linear(dims[-1], 3 * num_kp) # expression / delta
|
||||
|
||||
def _init_weights(self, m):
|
||||
if isinstance(m, (nn.Conv2d, nn.Linear)):
|
||||
trunc_normal_(m.weight, std=.02)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
def forward_features(self, x):
|
||||
for i in range(4):
|
||||
x = self.downsample_layers[i](x)
|
||||
x = self.stages[i](x)
|
||||
return self.norm(x.mean([-2, -1])) # global average pooling, (N, C, H, W) -> (N, C)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.forward_features(x)
|
||||
|
||||
# implicit keypoints
|
||||
kp = self.fc_kp(x)
|
||||
|
||||
# pose and expression deformation
|
||||
pitch = self.fc_pitch(x)
|
||||
yaw = self.fc_yaw(x)
|
||||
roll = self.fc_roll(x)
|
||||
t = self.fc_t(x)
|
||||
exp = self.fc_exp(x)
|
||||
scale = self.fc_scale(x)
|
||||
|
||||
ret_dct = {
|
||||
'pitch': pitch,
|
||||
'yaw': yaw,
|
||||
'roll': roll,
|
||||
't': t,
|
||||
'exp': exp,
|
||||
'scale': scale,
|
||||
|
||||
'kp': kp, # canonical keypoint
|
||||
}
|
||||
|
||||
return ret_dct
|
||||
|
||||
|
||||
def convnextv2_tiny(**kwargs):
|
||||
model = ConvNeXtV2(depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], **kwargs)
|
||||
return model
|
||||
@@ -0,0 +1,104 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
The module that predicting a dense motion from sparse motion representation given by kp_source and kp_driving
|
||||
"""
|
||||
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
import torch
|
||||
from .util import Hourglass, make_coordinate_grid, kp2gaussian
|
||||
|
||||
|
||||
class DenseMotionNetwork(nn.Module):
|
||||
def __init__(self, block_expansion, num_blocks, max_features, num_kp, feature_channel, reshape_depth, compress, estimate_occlusion_map=True):
|
||||
super(DenseMotionNetwork, self).__init__()
|
||||
self.hourglass = Hourglass(block_expansion=block_expansion, in_features=(num_kp+1)*(compress+1), max_features=max_features, num_blocks=num_blocks) # ~60+G
|
||||
|
||||
self.mask = nn.Conv3d(self.hourglass.out_filters, num_kp + 1, kernel_size=7, padding=3) # 65G! NOTE: computation cost is large
|
||||
self.compress = nn.Conv3d(feature_channel, compress, kernel_size=1) # 0.8G
|
||||
self.norm = nn.BatchNorm3d(compress, affine=True)
|
||||
self.num_kp = num_kp
|
||||
self.flag_estimate_occlusion_map = estimate_occlusion_map
|
||||
|
||||
if self.flag_estimate_occlusion_map:
|
||||
self.occlusion = nn.Conv2d(self.hourglass.out_filters*reshape_depth, 1, kernel_size=7, padding=3)
|
||||
else:
|
||||
self.occlusion = None
|
||||
|
||||
def create_sparse_motions(self, feature, kp_driving, kp_source):
|
||||
bs, _, d, h, w = feature.shape # (bs, 4, 16, 64, 64)
|
||||
identity_grid = make_coordinate_grid((d, h, w), ref=kp_source) # (16, 64, 64, 3)
|
||||
identity_grid = identity_grid.view(1, 1, d, h, w, 3) # (1, 1, d=16, h=64, w=64, 3)
|
||||
coordinate_grid = identity_grid - kp_driving.view(bs, self.num_kp, 1, 1, 1, 3)
|
||||
|
||||
k = coordinate_grid.shape[1]
|
||||
|
||||
# NOTE: there lacks an one-order flow
|
||||
driving_to_source = coordinate_grid + kp_source.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, 1+num_kp, d, h, w, 3)
|
||||
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, align_corners=False)
|
||||
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:] # (d=16, h=64, w=64)
|
||||
gaussian_driving = kp2gaussian(kp_driving, spatial_size=spatial_size, kp_variance=0.01) # (bs, num_kp, d, h, w)
|
||||
gaussian_source = kp2gaussian(kp_source, spatial_size=spatial_size, kp_variance=0.01) # (bs, num_kp, d, h, w)
|
||||
heatmap = gaussian_driving - gaussian_source # (bs, num_kp, d, h, w)
|
||||
|
||||
# adding background feature
|
||||
zeros = torch.zeros(heatmap.shape[0], 1, spatial_size[0], spatial_size[1], spatial_size[2]).type(heatmap.dtype).to(heatmap.device)
|
||||
heatmap = torch.cat([zeros, heatmap], dim=1)
|
||||
heatmap = heatmap.unsqueeze(2) # (bs, 1+num_kp, 1, d, h, w)
|
||||
return heatmap
|
||||
|
||||
def forward(self, feature, kp_driving, kp_source):
|
||||
bs, _, d, h, w = feature.shape # (bs, 32, 16, 64, 64)
|
||||
|
||||
feature = self.compress(feature) # (bs, 4, 16, 64, 64)
|
||||
feature = self.norm(feature) # (bs, 4, 16, 64, 64)
|
||||
feature = F.relu(feature) # (bs, 4, 16, 64, 64)
|
||||
|
||||
out_dict = dict()
|
||||
|
||||
# 1. deform 3d feature
|
||||
sparse_motion = self.create_sparse_motions(feature, kp_driving, kp_source) # (bs, 1+num_kp, d, h, w, 3)
|
||||
deformed_feature = self.create_deformed_feature(feature, sparse_motion) # (bs, 1+num_kp, c=4, d=16, h=64, w=64)
|
||||
|
||||
# 2. (bs, 1+num_kp, d, h, w)
|
||||
heatmap = self.create_heatmap_representations(deformed_feature, kp_driving, kp_source) # (bs, 1+num_kp, 1, d, h, w)
|
||||
|
||||
input = torch.cat([heatmap, deformed_feature], dim=2) # (bs, 1+num_kp, c=5, d=16, h=64, w=64)
|
||||
input = input.view(bs, -1, d, h, w) # (bs, (1+num_kp)*c=105, d=16, h=64, w=64)
|
||||
|
||||
prediction = self.hourglass(input)
|
||||
|
||||
mask = self.mask(prediction)
|
||||
mask = F.softmax(mask, dim=1) # (bs, 1+num_kp, d=16, h=64, w=64)
|
||||
out_dict['mask'] = mask
|
||||
mask = mask.unsqueeze(2) # (bs, num_kp+1, 1, d, h, w)
|
||||
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) mask take effect in this place
|
||||
deformation = deformation.permute(0, 2, 3, 4, 1) # (bs, d, h, w, 3)
|
||||
|
||||
out_dict['deformation'] = deformation
|
||||
|
||||
if self.flag_estimate_occlusion_map:
|
||||
bs, _, d, h, w = prediction.shape
|
||||
prediction_reshape = prediction.view(bs, -1, h, w)
|
||||
occlusion_map = torch.sigmoid(self.occlusion(prediction_reshape)) # Bx1x64x64
|
||||
out_dict['occlusion_map'] = occlusion_map
|
||||
|
||||
return out_dict
|
||||
@@ -0,0 +1,35 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
Motion extractor(M), which directly predicts the canonical keypoints, head pose and expression deformation of the input image
|
||||
"""
|
||||
|
||||
from torch import nn
|
||||
import torch
|
||||
|
||||
from .convnextv2 import convnextv2_tiny
|
||||
from .util import filter_state_dict
|
||||
|
||||
model_dict = {
|
||||
'convnextv2_tiny': convnextv2_tiny,
|
||||
}
|
||||
|
||||
|
||||
class MotionExtractor(nn.Module):
|
||||
def __init__(self, **kwargs):
|
||||
super(MotionExtractor, self).__init__()
|
||||
|
||||
# default is convnextv2_base
|
||||
backbone = kwargs.get('backbone', 'convnextv2_tiny')
|
||||
self.detector = model_dict.get(backbone)(**kwargs)
|
||||
|
||||
def load_pretrained(self, init_path: str):
|
||||
if init_path not in (None, ''):
|
||||
state_dict = torch.load(init_path, map_location=lambda storage, loc: storage)['model']
|
||||
state_dict = filter_state_dict(state_dict, remove_name='head')
|
||||
ret = self.detector.load_state_dict(state_dict, strict=False)
|
||||
print(f'Load pretrained model from {init_path}, ret: {ret}')
|
||||
|
||||
def forward(self, x):
|
||||
out = self.detector(x)
|
||||
return out
|
||||
@@ -0,0 +1,59 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
Spade decoder(G) defined in the paper, which input the warped feature to generate the animated image.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
from .util import SPADEResnetBlock
|
||||
|
||||
|
||||
class SPADEDecoder(nn.Module):
|
||||
def __init__(self, upscale=1, max_features=256, block_expansion=64, out_channels=64, num_down_blocks=2):
|
||||
for i in range(num_down_blocks):
|
||||
input_channels = min(max_features, block_expansion * (2 ** (i + 1)))
|
||||
self.upscale = upscale
|
||||
super().__init__()
|
||||
norm_G = 'spadespectralinstance'
|
||||
label_num_channels = input_channels # 256
|
||||
|
||||
self.fc = nn.Conv2d(input_channels, 2 * input_channels, 3, padding=1)
|
||||
self.G_middle_0 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels)
|
||||
self.G_middle_1 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels)
|
||||
self.G_middle_2 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels)
|
||||
self.G_middle_3 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels)
|
||||
self.G_middle_4 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels)
|
||||
self.G_middle_5 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels)
|
||||
self.up_0 = SPADEResnetBlock(2 * input_channels, input_channels, norm_G, label_num_channels)
|
||||
self.up_1 = SPADEResnetBlock(input_channels, out_channels, norm_G, label_num_channels)
|
||||
self.up = nn.Upsample(scale_factor=2)
|
||||
|
||||
if self.upscale is None or self.upscale <= 1:
|
||||
self.conv_img = nn.Conv2d(out_channels, 3, 3, padding=1)
|
||||
else:
|
||||
self.conv_img = nn.Sequential(
|
||||
nn.Conv2d(out_channels, 3 * (2 * 2), kernel_size=3, padding=1),
|
||||
nn.PixelShuffle(upscale_factor=2)
|
||||
)
|
||||
|
||||
def forward(self, feature):
|
||||
seg = feature # Bx256x64x64
|
||||
x = self.fc(feature) # Bx512x64x64
|
||||
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) # Bx512x64x64 -> Bx512x128x128
|
||||
x = self.up_0(x, seg) # Bx512x128x128 -> Bx256x128x128
|
||||
x = self.up(x) # Bx256x128x128 -> Bx256x256x256
|
||||
x = self.up_1(x, seg) # Bx256x256x256 -> Bx64x256x256
|
||||
|
||||
x = self.conv_img(F.leaky_relu(x, 2e-1)) # Bx64x256x256 -> Bx3xHxW
|
||||
x = torch.sigmoid(x) # Bx3xHxW
|
||||
|
||||
return x
|
||||
@@ -0,0 +1,38 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
Stitching module(S) and two retargeting modules(R) defined in the paper.
|
||||
|
||||
- The stitching module pastes the animated portrait back into the original image space without pixel misalignment, such as in
|
||||
the stitching region.
|
||||
|
||||
- The eyes retargeting module is designed to address the issue of incomplete eye closure during cross-id reenactment, especially
|
||||
when a person with small eyes drives a person with larger eyes.
|
||||
|
||||
- The lip retargeting module is designed similarly to the eye retargeting module, and can also normalize the input by ensuring that
|
||||
the lips are in a closed state, which facilitates better animation driving.
|
||||
"""
|
||||
from torch import nn
|
||||
|
||||
|
||||
class StitchingRetargetingNetwork(nn.Module):
|
||||
def __init__(self, input_size, hidden_sizes, output_size):
|
||||
super(StitchingRetargetingNetwork, self).__init__()
|
||||
layers = []
|
||||
for i in range(len(hidden_sizes)):
|
||||
if i == 0:
|
||||
layers.append(nn.Linear(input_size, hidden_sizes[i]))
|
||||
else:
|
||||
layers.append(nn.Linear(hidden_sizes[i - 1], hidden_sizes[i]))
|
||||
layers.append(nn.ReLU(inplace=True))
|
||||
layers.append(nn.Linear(hidden_sizes[-1], output_size))
|
||||
self.mlp = nn.Sequential(*layers)
|
||||
|
||||
def initialize_weights_to_zero(self):
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Linear):
|
||||
nn.init.zeros_(m.weight)
|
||||
nn.init.zeros_(m.bias)
|
||||
|
||||
def forward(self, x):
|
||||
return self.mlp(x)
|
||||
@@ -0,0 +1,452 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
This file defines various neural network modules and utility functions, including convolutional and residual blocks,
|
||||
normalizations, and functions for spatial transformation and tensor manipulation.
|
||||
"""
|
||||
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
import torch
|
||||
import torch.nn.utils.spectral_norm as spectral_norm
|
||||
import math
|
||||
import warnings
|
||||
import collections.abc
|
||||
from itertools import repeat
|
||||
|
||||
def kp2gaussian(kp, spatial_size, kp_variance):
|
||||
"""
|
||||
Transform a keypoint into gaussian like representation
|
||||
"""
|
||||
mean = kp
|
||||
|
||||
coordinate_grid = make_coordinate_grid(spatial_size, mean)
|
||||
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(spatial_size, ref, **kwargs):
|
||||
d, h, w = spatial_size
|
||||
x = torch.arange(w).type(ref.dtype).to(ref.device)
|
||||
y = torch.arange(h).type(ref.dtype).to(ref.device)
|
||||
z = torch.arange(d).type(ref.dtype).to(ref.device)
|
||||
|
||||
# NOTE: must be right-down-in
|
||||
x = (2 * (x / (w - 1)) - 1) # the x axis faces to the right
|
||||
y = (2 * (y / (h - 1)) - 1) # the y axis faces to the bottom
|
||||
z = (2 * (z / (d - 1)) - 1) # the z axis faces to the inner
|
||||
|
||||
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 ConvT2d(nn.Module):
|
||||
"""
|
||||
Upsampling block for use in decoder.
|
||||
"""
|
||||
|
||||
def __init__(self, in_features, out_features, kernel_size=3, stride=2, padding=1, output_padding=1):
|
||||
super(ConvT2d, self).__init__()
|
||||
|
||||
self.convT = nn.ConvTranspose2d(in_features, out_features, kernel_size=kernel_size, stride=stride,
|
||||
padding=padding, output_padding=output_padding)
|
||||
self.norm = nn.InstanceNorm2d(out_features)
|
||||
|
||||
def forward(self, x):
|
||||
out = self.convT(x)
|
||||
out = self.norm(out)
|
||||
out = F.leaky_relu(out)
|
||||
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 = nn.BatchNorm3d(in_features, affine=True)
|
||||
self.norm2 = nn.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 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 = nn.BatchNorm3d(out_features, affine=True)
|
||||
|
||||
def forward(self, x):
|
||||
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 = nn.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 = nn.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 = nn.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 + in_features
|
||||
|
||||
self.conv = nn.Conv3d(in_channels=self.out_filters, out_channels=self.out_filters, kernel_size=3, padding=1)
|
||||
self.norm = nn.BatchNorm3d(self.out_filters, affine=True)
|
||||
|
||||
def forward(self, x):
|
||||
out = x.pop()
|
||||
for up_block in self.up_blocks:
|
||||
out = up_block(out)
|
||||
skip = x.pop()
|
||||
out = torch.cat([out, skip], dim=1)
|
||||
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 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)
|
||||
|
||||
|
||||
def filter_state_dict(state_dict, remove_name='fc'):
|
||||
new_state_dict = {}
|
||||
for key in state_dict:
|
||||
if remove_name in key:
|
||||
continue
|
||||
new_state_dict[key] = state_dict[key]
|
||||
return new_state_dict
|
||||
|
||||
|
||||
class GRN(nn.Module):
|
||||
""" GRN (Global Response Normalization) layer
|
||||
"""
|
||||
|
||||
def __init__(self, dim):
|
||||
super().__init__()
|
||||
self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim))
|
||||
self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim))
|
||||
|
||||
def forward(self, x):
|
||||
Gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True)
|
||||
Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
|
||||
return self.gamma * (x * Nx) + self.beta + x
|
||||
|
||||
|
||||
class LayerNorm(nn.Module):
|
||||
r""" LayerNorm that supports two data formats: channels_last (default) or channels_first.
|
||||
The ordering of the dimensions in the inputs. channels_last corresponds to inputs with
|
||||
shape (batch_size, height, width, channels) while channels_first corresponds to inputs
|
||||
with shape (batch_size, channels, height, width).
|
||||
"""
|
||||
|
||||
def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.ones(normalized_shape))
|
||||
self.bias = nn.Parameter(torch.zeros(normalized_shape))
|
||||
self.eps = eps
|
||||
self.data_format = data_format
|
||||
if self.data_format not in ["channels_last", "channels_first"]:
|
||||
raise NotImplementedError
|
||||
self.normalized_shape = (normalized_shape, )
|
||||
|
||||
def forward(self, x):
|
||||
if self.data_format == "channels_last":
|
||||
return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
|
||||
elif self.data_format == "channels_first":
|
||||
u = x.mean(1, keepdim=True)
|
||||
s = (x - u).pow(2).mean(1, keepdim=True)
|
||||
x = (x - u) / torch.sqrt(s + self.eps)
|
||||
x = self.weight[:, None, None] * x + self.bias[:, None, None]
|
||||
return x
|
||||
|
||||
|
||||
def _no_grad_trunc_normal_(tensor, mean, std, a, b):
|
||||
# Cut & paste from PyTorch official master until it's in a few official releases - RW
|
||||
# Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
|
||||
def norm_cdf(x):
|
||||
# Computes standard normal cumulative distribution function
|
||||
return (1. + math.erf(x / math.sqrt(2.))) / 2.
|
||||
|
||||
if (mean < a - 2 * std) or (mean > b + 2 * std):
|
||||
warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
|
||||
"The distribution of values may be incorrect.",
|
||||
stacklevel=2)
|
||||
|
||||
with torch.no_grad():
|
||||
# Values are generated by using a truncated uniform distribution and
|
||||
# then using the inverse CDF for the normal distribution.
|
||||
# Get upper and lower cdf values
|
||||
l = norm_cdf((a - mean) / std)
|
||||
u = norm_cdf((b - mean) / std)
|
||||
|
||||
# Uniformly fill tensor with values from [l, u], then translate to
|
||||
# [2l-1, 2u-1].
|
||||
tensor.uniform_(2 * l - 1, 2 * u - 1)
|
||||
|
||||
# Use inverse cdf transform for normal distribution to get truncated
|
||||
# standard normal
|
||||
tensor.erfinv_()
|
||||
|
||||
# Transform to proper mean, std
|
||||
tensor.mul_(std * math.sqrt(2.))
|
||||
tensor.add_(mean)
|
||||
|
||||
# Clamp to ensure it's in the proper range
|
||||
tensor.clamp_(min=a, max=b)
|
||||
return tensor
|
||||
|
||||
|
||||
def drop_path(x, drop_prob=0., training=False, scale_by_keep=True):
|
||||
""" Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
|
||||
|
||||
This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
|
||||
the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
|
||||
See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
|
||||
changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
|
||||
'survival rate' as the argument.
|
||||
|
||||
"""
|
||||
if drop_prob == 0. or not training:
|
||||
return x
|
||||
keep_prob = 1 - drop_prob
|
||||
shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
|
||||
random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
|
||||
if keep_prob > 0.0 and scale_by_keep:
|
||||
random_tensor.div_(keep_prob)
|
||||
return x * random_tensor
|
||||
|
||||
|
||||
class DropPath(nn.Module):
|
||||
""" Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
|
||||
"""
|
||||
|
||||
def __init__(self, drop_prob=None, scale_by_keep=True):
|
||||
super(DropPath, self).__init__()
|
||||
self.drop_prob = drop_prob
|
||||
self.scale_by_keep = scale_by_keep
|
||||
|
||||
def forward(self, x):
|
||||
return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)
|
||||
|
||||
|
||||
def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
|
||||
return _no_grad_trunc_normal_(tensor, mean, std, a, b)
|
||||
|
||||
# From PyTorch internals
|
||||
def _ntuple(n):
|
||||
def parse(x):
|
||||
if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
|
||||
return tuple(x)
|
||||
return tuple(repeat(x, n))
|
||||
return parse
|
||||
|
||||
to_2tuple = _ntuple(2)
|
||||
@@ -0,0 +1,77 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
Warping field estimator(W) defined in the paper, which generates a warping field using the implicit
|
||||
keypoint representations x_s and x_d, and employs this flow field to warp the source feature volume f_s.
|
||||
"""
|
||||
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
from .util import SameBlock2d
|
||||
from .dense_motion import DenseMotionNetwork
|
||||
|
||||
|
||||
class WarpingNetwork(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
num_kp,
|
||||
block_expansion,
|
||||
max_features,
|
||||
num_down_blocks,
|
||||
reshape_channel,
|
||||
estimate_occlusion_map=False,
|
||||
dense_motion_params=None,
|
||||
**kwargs
|
||||
):
|
||||
super(WarpingNetwork, self).__init__()
|
||||
|
||||
self.upscale = kwargs.get('upscale', 1)
|
||||
self.flag_use_occlusion_map = kwargs.get('flag_use_occlusion_map', True)
|
||||
|
||||
if dense_motion_params is not None:
|
||||
self.dense_motion_network = DenseMotionNetwork(
|
||||
num_kp=num_kp,
|
||||
feature_channel=reshape_channel,
|
||||
estimate_occlusion_map=estimate_occlusion_map,
|
||||
**dense_motion_params
|
||||
)
|
||||
else:
|
||||
self.dense_motion_network = None
|
||||
|
||||
self.third = SameBlock2d(max_features, block_expansion * (2 ** num_down_blocks), kernel_size=(3, 3), padding=(1, 1), lrelu=True)
|
||||
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)
|
||||
|
||||
self.estimate_occlusion_map = estimate_occlusion_map
|
||||
|
||||
def deform_input(self, inp, deformation):
|
||||
return F.grid_sample(inp, deformation, align_corners=False)
|
||||
|
||||
def forward(self, feature_3d, kp_driving, kp_source):
|
||||
if self.dense_motion_network is not None:
|
||||
# Feature warper, Transforming feature representation according to deformation and occlusion
|
||||
dense_motion = self.dense_motion_network(
|
||||
feature=feature_3d, kp_driving=kp_driving, kp_source=kp_source
|
||||
)
|
||||
if 'occlusion_map' in dense_motion:
|
||||
occlusion_map = dense_motion['occlusion_map'] # Bx1x64x64
|
||||
else:
|
||||
occlusion_map = None
|
||||
|
||||
deformation = dense_motion['deformation'] # Bx16x64x64x3
|
||||
out = self.deform_input(feature_3d, deformation) # Bx32x16x64x64
|
||||
|
||||
bs, c, d, h, w = out.shape # Bx32x16x64x64
|
||||
out = out.view(bs, c * d, h, w) # -> Bx512x64x64
|
||||
out = self.third(out) # -> Bx256x64x64
|
||||
out = self.fourth(out) # -> Bx256x64x64
|
||||
|
||||
if self.flag_use_occlusion_map and (occlusion_map is not None):
|
||||
out = out * occlusion_map
|
||||
|
||||
ret_dct = {
|
||||
'occlusion_map': occlusion_map,
|
||||
'deformation': deformation,
|
||||
'out': out,
|
||||
}
|
||||
|
||||
return ret_dct
|
||||
Reference in New Issue
Block a user