Spaces:
Runtime error
Runtime error
import torch.nn as nn | |
import torch | |
import numpy as np | |
from .skeleton_DME import SkeletonConv, SkeletonPool, find_neighbor, build_edge_topology | |
from .skeleton import SkeletonResidual | |
from .decoders import VQDecoderV3 | |
class LocalEncoder(nn.Module): | |
def __init__(self, args, topology): | |
super(LocalEncoder, self).__init__() | |
args.channel_base = 6 | |
args.activation = "tanh" | |
args.use_residual_blocks = True | |
args.z_dim = 1024 | |
args.temporal_scale = 8 | |
args.kernel_size = 4 | |
args.num_layers = args.vae_layer | |
args.skeleton_dist = 2 | |
args.extra_conv = 0 | |
# check how to reflect in 1d | |
args.padding_mode = "constant" | |
args.skeleton_pool = "mean" | |
args.upsampling = "linear" | |
self.topologies = [topology] | |
self.channel_base = [args.channel_base] | |
self.channel_list = [] | |
self.edge_num = [len(topology)] | |
self.pooling_list = [] | |
self.layers = nn.ModuleList() | |
self.args = args | |
# self.convs = [] | |
kernel_size = args.kernel_size | |
kernel_even = False if kernel_size % 2 else True | |
padding = (kernel_size - 1) // 2 | |
bias = True | |
self.grow = args.vae_grow | |
for i in range(args.num_layers): | |
self.channel_base.append(self.channel_base[-1] * self.grow[i]) | |
for i in range(args.num_layers): | |
seq = [] | |
neighbour_list = find_neighbor(self.topologies[i], args.skeleton_dist) | |
in_channels = self.channel_base[i] * self.edge_num[i] | |
out_channels = self.channel_base[i + 1] * self.edge_num[i] | |
if i == 0: | |
self.channel_list.append(in_channels) | |
self.channel_list.append(out_channels) | |
last_pool = True if i == args.num_layers - 1 else False | |
# (T, J, D) => (T, J', D) | |
pool = SkeletonPool( | |
edges=self.topologies[i], | |
pooling_mode=args.skeleton_pool, | |
channels_per_edge=out_channels // len(neighbour_list), | |
last_pool=last_pool, | |
) | |
if args.use_residual_blocks: | |
# (T, J, D) => (T/2, J', 2D) | |
seq.append( | |
SkeletonResidual( | |
self.topologies[i], | |
neighbour_list, | |
joint_num=self.edge_num[i], | |
in_channels=in_channels, | |
out_channels=out_channels, | |
kernel_size=kernel_size, | |
stride=2, | |
padding=padding, | |
padding_mode=args.padding_mode, | |
bias=bias, | |
extra_conv=args.extra_conv, | |
pooling_mode=args.skeleton_pool, | |
activation=args.activation, | |
last_pool=last_pool, | |
) | |
) | |
else: | |
for _ in range(args.extra_conv): | |
# (T, J, D) => (T, J, D) | |
seq.append( | |
SkeletonConv( | |
neighbour_list, | |
in_channels=in_channels, | |
out_channels=in_channels, | |
joint_num=self.edge_num[i], | |
kernel_size=kernel_size - 1 if kernel_even else kernel_size, | |
stride=1, | |
padding=padding, | |
padding_mode=args.padding_mode, | |
bias=bias, | |
) | |
) | |
seq.append(nn.PReLU() if args.activation == "relu" else nn.Tanh()) | |
# (T, J, D) => (T/2, J, 2D) | |
seq.append( | |
SkeletonConv( | |
neighbour_list, | |
in_channels=in_channels, | |
out_channels=out_channels, | |
joint_num=self.edge_num[i], | |
kernel_size=kernel_size, | |
stride=2, | |
padding=padding, | |
padding_mode=args.padding_mode, | |
bias=bias, | |
add_offset=False, | |
in_offset_channel=3 * self.channel_base[i] // self.channel_base[0], | |
) | |
) | |
# self.convs.append(seq[-1]) | |
seq.append(pool) | |
seq.append(nn.PReLU() if args.activation == "relu" else nn.Tanh()) | |
self.layers.append(nn.Sequential(*seq)) | |
self.topologies.append(pool.new_edges) | |
self.pooling_list.append(pool.pooling_list) | |
self.edge_num.append(len(self.topologies[-1])) | |
# in_features = self.channel_base[-1] * len(self.pooling_list[-1]) | |
# in_features *= int(args.temporal_scale / 2) | |
# self.reduce = nn.Linear(in_features, args.z_dim) | |
# self.mu = nn.Linear(in_features, args.z_dim) | |
# self.logvar = nn.Linear(in_features, args.z_dim) | |
def forward(self, input): | |
# bs, n, c = input.shape[0], input.shape[1], input.shape[2] | |
output = input.permute(0, 2, 1) # input.reshape(bs, n, -1, 6) | |
for layer in self.layers: | |
output = layer(output) | |
# output = output.view(output.shape[0], -1) | |
output = output.permute(0, 2, 1) | |
return output | |
def reparameterize(mu, logvar): | |
std = torch.exp(0.5 * logvar) | |
eps = torch.randn_like(std) | |
return mu + eps * std | |
class VAEConv(nn.Module): | |
def __init__(self, args): | |
super(VAEConv, self).__init__() | |
# self.encoder = VQEncoderV3(args) | |
# self.decoder = VQDecoderV3(args) | |
self.fc_mu = nn.Linear(args.vae_length, args.vae_length) | |
self.fc_logvar = nn.Linear(args.vae_length, args.vae_length) | |
self.variational = args.variational | |
def forward(self, inputs): | |
pre_latent = self.encoder(inputs) | |
mu, logvar = None, None | |
if self.variational: | |
mu = self.fc_mu(pre_latent) | |
logvar = self.fc_logvar(pre_latent) | |
pre_latent = reparameterize(mu, logvar) | |
rec_pose = self.decoder(pre_latent) | |
return { | |
"poses_feat": pre_latent, | |
"rec_pose": rec_pose, | |
"pose_mu": mu, | |
"pose_logvar": logvar, | |
} | |
def map2latent(self, inputs): | |
pre_latent = self.encoder(inputs) | |
if self.variational: | |
mu = self.fc_mu(pre_latent) | |
logvar = self.fc_logvar(pre_latent) | |
pre_latent = reparameterize(mu, logvar) | |
return pre_latent | |
def decode(self, pre_latent): | |
rec_pose = self.decoder(pre_latent) | |
return rec_pose | |
class VAESKConv(VAEConv): | |
def __init__(self, args, model_save_path="./emage/"): | |
# args = args() | |
super(VAESKConv, self).__init__(args) | |
smpl_fname = model_save_path + "smplx_models/smplx/SMPLX_NEUTRAL_2020.npz" | |
smpl_data = np.load(smpl_fname, encoding="latin1") | |
parents = smpl_data["kintree_table"][0].astype(np.int32) | |
edges = build_edge_topology(parents) | |
self.encoder = LocalEncoder(args, edges) | |
self.decoder = VQDecoderV3(args) | |