Spaces:
Runtime error
Runtime error
File size: 1,585 Bytes
a22eb82 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 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} |