from torch import nn import torch.nn.functional as F import torch import numpy as np import cv2 from FONT.sync_batchnorm import SynchronizedBatchNorm2d as BatchNorm2d from .stylegan2 import Generator import torch.nn as nn import math import torch.utils.model_zoo as model_zoo from .function import adaptive_instance_normalization as adain import pdb from .conv import Conv2dTranspose, Conv2d, nonorm_Conv2d # Misc img2mse = lambda x, y : torch.mean((x - y) ** 2) mse2psnr = lambda x : -10. * torch.log(x) / torch.log(torch.Tensor([10.])) to8b = lambda x : (255*np.clip(x,0,1)).astype(np.uint8) class InstanceNorm(nn.Module): def __init__(self, epsilon=1e-8): """ @notice: avoid in-place ops. https://discuss.pytorch.org/t/encounter-the-runtimeerror-one-of-the-variables-needed-for-gradient-computation-has-been-modified-by-an-inplace-operation/836/3 """ super(InstanceNorm, self).__init__() self.epsilon = epsilon def forward(self, x): x = x - torch.mean(x, (2, 3), True) tmp = torch.mul(x, x) # or x ** 2 tmp = torch.rsqrt(torch.mean(tmp, (2, 3), True) + self.epsilon) return x * tmp class ApplyStyle(nn.Module): """ @ref: https://github.com/lernapparat/lernapparat/blob/master/style_gan/pytorch_style_gan.ipynb """ def __init__(self, latent_size, channels, use_wscale): super(ApplyStyle, self).__init__() self.linear = FC(latent_size, channels * 2, gain=1.0, use_wscale=use_wscale) def forward(self, x, latent): style = self.linear(latent) # style => [batch_size, n_channels*2] shape = [-1, 2, x.size(1), 1, 1] style = style.view(shape) # [batch_size, 2, n_channels, ...] x = x * (style[:, 0] + 1.) + style[:, 1] return x class FC(nn.Module): def __init__(self, in_channels, out_channels, gain=2**(0.5), use_wscale=False, lrmul=1.0, bias=True): """ The complete conversion of Dense/FC/Linear Layer of original Tensorflow version. """ super(FC, self).__init__() he_std = gain * in_channels ** (-0.5) # He init if use_wscale: init_std = 1.0 / lrmul self.w_lrmul = he_std * lrmul else: init_std = he_std / lrmul self.w_lrmul = lrmul self.weight = torch.nn.Parameter(torch.randn(out_channels, in_channels) * init_std) if bias: self.bias = torch.nn.Parameter(torch.zeros(out_channels)) self.b_lrmul = lrmul else: self.bias = None def forward(self, x): if self.bias is not None: out = F.linear(x, self.weight * self.w_lrmul, self.bias * self.b_lrmul) else: out = F.linear(x, self.weight * self.w_lrmul) out = F.leaky_relu(out, 0.2, inplace=True) return out # Positional encoding (section 5.1) class Embedder: def __init__(self, **kwargs): self.kwargs = kwargs self.create_embedding_fn() def create_embedding_fn(self): embed_fns = [] d = self.kwargs['input_dims'] out_dim = 0 if self.kwargs['include_input']: embed_fns.append(lambda x : x) out_dim += d max_freq = self.kwargs['max_freq_log2'] N_freqs = self.kwargs['num_freqs'] if self.kwargs['log_sampling']: freq_bands = 2.**torch.linspace(0., max_freq, steps=N_freqs) else: freq_bands = torch.linspace(2.**0., 2.**max_freq, steps=N_freqs) for freq in freq_bands: for p_fn in self.kwargs['periodic_fns']: embed_fns.append(lambda x, p_fn=p_fn, freq=freq : p_fn(x * freq)) out_dim += d self.embed_fns = embed_fns self.out_dim = out_dim def embed(self, inputs): return torch.cat([fn(inputs) for fn in self.embed_fns], -1) def get_embedder(multires, i=0): if i == -1: return nn.Identity(), 6 embed_kwargs = { 'include_input' : True, 'input_dims' : 6, 'max_freq_log2' : multires-1, 'num_freqs' : multires, 'log_sampling' : True, 'periodic_fns' : [torch.sin, torch.cos], } embedder_obj = Embedder(**embed_kwargs) embed = lambda x, eo=embedder_obj : eo.embed(x) return embed, embedder_obj.out_dim def draw_heatmap(landmark, width, height): batch = landmark.shape[0] number = landmark.shape[1] heatmap = np.zeros((batch, number,width, height), dtype=np.float32) # draw mouth from mouth landmarks, landmarks: mouth landmark points, format: x1, y1, x2, y2, ..., x20, landmark = (landmark+1)*29 for i in range(batch): for pts_idx in range(number): if int(landmark[i,pts_idx,0])<0: landmark[i,pts_idx,0] = 0 if int(landmark[i,pts_idx,1])<0: landmark[i,pts_idx,1] = 0 if int(landmark[i,pts_idx,0])>57: landmark[i,pts_idx,0] = 57 if int(landmark[i,pts_idx,1])>57: landmark[i,pts_idx,1] = 57 heatmap[i,pts_idx, int(landmark[i,pts_idx,1]), int(landmark[i,pts_idx,0])]=1 if heatmap[i,pts_idx].sum()== 1 : heatmap[i,pts_idx] = cv2.GaussianBlur(heatmap[i,pts_idx], ksize=(3, 3), sigmaX=1, sigmaY=1) heatmap = torch.tensor(heatmap).cuda() return heatmap class NA_net(nn.Module): def __init__(self): super(NA_net, self).__init__() self.decon = nn.Sequential( nn.ConvTranspose2d(1, 16, kernel_size=(2,3), stride=2, padding=(2,1), bias=True),#16,16 nn.BatchNorm2d(16), nn.ReLU(True), nn.ConvTranspose2d(16, 32, kernel_size=4, stride=2, padding=1, bias=True),#8,8 nn.BatchNorm2d(32), nn.ReLU(True), nn.ConvTranspose2d(32, 32+3, kernel_size=4, stride=2, padding=1, bias=True)#16,16 ) def forward(self, neutral): feature = neutral.unsqueeze(1) current_feature = self.decon(feature) return current_feature class AT_net(nn.Module): def __init__(self): super(AT_net, self).__init__() down_blocks = [] for i in range(8): down_blocks.append(DownBlock2d(3 if i == 0 else 2 * (2 ** i), 2 * (2 ** (i + 1)), kernel_size=3, padding=1)) self.down_blocks = nn.ModuleList(down_blocks) # self.lmark_encoder = nn.Sequential( # nn.Linear(16,256), # nn.ReLU(True), # nn.Linear(256,512), # nn.ReLU(True), # ) self.pose_encoder = nn.Sequential( nn.Linear(6,128), nn.ReLU(True), nn.Linear(128,256), nn.ReLU(True), ) self.audio_eocder = nn.Sequential( conv2d(1,64,3,1,1), conv2d(64,128,3,1,1), nn.MaxPool2d(3, stride=(1,2)), conv2d(128,256,3,1,1), conv2d(256,256,3,1,1), conv2d(256,512,3,1,1), nn.MaxPool2d(3, stride=(2,2)) ) self.audio_eocder_fc = nn.Sequential( nn.Linear(1024 *12,2048), nn.ReLU(True), nn.Linear(2048,256), nn.ReLU(True), ) self.lstm = nn.LSTM(256*4,256,3,batch_first = True) # self.lstm_fc = nn.Sequential( # nn.Linear(256,16), # ) self.decon = nn.Sequential( nn.ConvTranspose2d(256, 256, kernel_size=6, stride=2, padding=1, bias=True),#4,4 nn.BatchNorm2d(256), nn.ReLU(True), nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1, bias=True),#8,8 nn.BatchNorm2d(128), nn.ReLU(True), nn.ConvTranspose2d(128, 128, kernel_size=4, stride=2, padding=1, bias=True), #16,16 nn.BatchNorm2d(128), nn.ReLU(True), nn.ConvTranspose2d(128, 128, kernel_size=4, stride=2, padding=1, bias=True),#32,32 nn.BatchNorm2d(128), nn.ReLU(True), nn.ConvTranspose2d(128, 32+3, kernel_size=4, stride=2, padding=1, bias=True),#64,64 # nn.ConvTranspose2d(128, 32*4, kernel_size=2, stride=2, padding=3, bias=True),#64,64 ) self.generator = Generator(64,256,8) def forward(self, example_image, audio, pose, jaco_net): hidden = ( torch.autograd.Variable(torch.zeros(3, audio.size(0), 256).cuda()), torch.autograd.Variable(torch.zeros(3, audio.size(0), 256).cuda())) outs = example_image for down_block in self.down_blocks: outs = down_block(outs) image_feature = outs image_feature = image_feature.view(image_feature.shape[0], -1) lstm_input = [] for step_t in range(audio.size(1)): current_audio = audio[ : ,step_t , :, :].unsqueeze(1) current_feature = self.audio_eocder(current_audio) current_feature = current_feature.view(current_feature.size(0), -1) current_feature = self.audio_eocder_fc(current_feature) pose_f = self.pose_encoder(pose[:,step_t]) features = torch.cat([image_feature, current_feature, pose_f], 1) lstm_input.append(features) lstm_input = torch.stack(lstm_input, dim = 1) lstm_out, _ = self.lstm(lstm_input, hidden) fc_out = [] deco_out = [] for step_t in range(audio.size(1)): fc_in = lstm_out[:,step_t,:] # fc_out.append(self.lstm_fc(fc_in)) if jaco_net == 'cnn': fc_feature = torch.unsqueeze(fc_in,2) fc_feature = torch.unsqueeze(fc_feature,3) deco_out.append(self.decon(fc_feature)) elif jaco_net == 'gan': result,_ = self.generator([fc_in]) deco_out.append(result) else: raise Exception("jaco_net type wrong") return torch.stack(deco_out,dim=1) class AT_net2(nn.Module): def __init__(self, device): super(AT_net2, self).__init__() self.device = device down_blocks = [] for i in range(8): down_blocks.append(DownBlock2d(3 if i == 0 else 2 * (2 ** i), 2 * (2 ** (i + 1)), kernel_size=3, padding=1)) self.down_blocks = nn.ModuleList(down_blocks) # self.lmark_encoder = nn.Sequential( # nn.Linear(16,256), # nn.ReLU(True), # nn.Linear(256,512), # nn.ReLU(True), # ) self.pose_encoder = nn.Sequential( nn.Linear(6,128), nn.ReLU(True), nn.Linear(128,256), nn.ReLU(True), ) self.audio_eocder = nn.Sequential( conv2d(1,64,3,1,1), conv2d(64,128,3,1,1), nn.MaxPool2d(3, stride=(1,2)), conv2d(128,256,3,1,1), conv2d(256,256,3,1,1), conv2d(256,512,3,1,1), nn.MaxPool2d(3, stride=(2,2)) ) self.audio_eocder_fc = nn.Sequential( nn.Linear(1024 *12,2048), nn.ReLU(True), nn.Linear(2048,256), nn.ReLU(True), ) self.lstm = nn.LSTM(256*4,256,3,batch_first = True) # self.lstm_fc = nn.Sequential( # nn.Linear(256,16), # ) self.decon = nn.Sequential( nn.ConvTranspose2d(256, 256, kernel_size=6, stride=2, padding=1, bias=True),#4,4 nn.BatchNorm2d(256), nn.ReLU(True), nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1, bias=True),#8,8 nn.BatchNorm2d(128), nn.ReLU(True), nn.ConvTranspose2d(128, 128, kernel_size=4, stride=2, padding=1, bias=True), #16,16 nn.BatchNorm2d(128), nn.ReLU(True), nn.ConvTranspose2d(128, 128, kernel_size=4, stride=2, padding=1, bias=True),#32,32 nn.BatchNorm2d(128), nn.ReLU(True), nn.ConvTranspose2d(128, 32+3, kernel_size=4, stride=2, padding=1, bias=True),#64,64 # nn.ConvTranspose2d(128, 32*4, kernel_size=2, stride=2, padding=3, bias=True),#64,64 ) self.generator = Generator(64,256,8) def forward(self, example_image, audio, pose, jaco_net, weight): hidden_ele1 = torch.zeros(3, audio.size(0), 256) hidden_ele2 = torch.zeros(3, audio.size(0), 256) if self.device == 'cuda': hidden_ele1 = hidden_ele1.cuda() hidden_ele2 = hidden_ele2.cuda() hidden = (torch.autograd.Variable(hidden_ele1), torch.autograd.Variable(hidden_ele2) ) outs = example_image for down_block in self.down_blocks: outs = down_block(outs) image_feature = outs image_feature = image_feature.view(image_feature.shape[0], -1) lstm_input = [] for step_t in range(audio.size(1)): current_audio = audio[ : ,step_t , :, :].unsqueeze(1) current_feature = self.audio_eocder(current_audio) current_feature = current_feature.view(current_feature.size(0), -1) current_feature = self.audio_eocder_fc(current_feature)*weight pose_f = self.pose_encoder(pose[:,step_t]) features = torch.cat([image_feature, current_feature, pose_f], 1) lstm_input.append(features) lstm_input = torch.stack(lstm_input, dim = 1) lstm_out, _ = self.lstm(lstm_input, hidden) fc_out = [] deco_out = [] for step_t in range(audio.size(1)): fc_in = lstm_out[:,step_t,:] # fc_out.append(self.lstm_fc(fc_in)) if jaco_net == 'cnn': fc_feature = torch.unsqueeze(fc_in,2) fc_feature = torch.unsqueeze(fc_feature,3) deco_out.append(self.decon(fc_feature)) elif jaco_net == 'gan': result,_ = self.generator([fc_in]) deco_out.append(result) else: raise Exception("jaco_net type wrong") return torch.stack(deco_out,dim=1) class AT_net_liujin(nn.Module): def __init__(self): super(AT_net_liujin, self).__init__() down_blocks = [] for i in range(8): down_blocks.append(DownBlock2d(3 if i == 0 else 2 * (2 ** i), 2 * (2 ** (i + 1)), kernel_size=3, padding=1)) self.down_blocks = nn.ModuleList(down_blocks) # self.lmark_encoder = nn.Sequential( # nn.Linear(16,256), # nn.ReLU(True), # nn.Linear(256,512), # nn.ReLU(True), # ) self.pose_encoder = nn.Sequential( nn.Linear(6,128), nn.ReLU(True), nn.Linear(128,256), nn.ReLU(True), ) self.audio_eocder = nn.Sequential( conv2d(1,64,3,1,1), conv2d(64,128,3,1,1), nn.MaxPool2d(3, stride=(1,2)), conv2d(128,256,3,1,1), conv2d(256,256,3,1,1), conv2d(256,512,3,1,1), nn.MaxPool2d(3, stride=(2,2)) ) self.audio_eocder_fc = nn.Sequential( nn.Linear(1024 *12,2048), nn.ReLU(True), nn.Linear(2048,256), nn.ReLU(True), ) self.audio_encoder_liujin = nn.Sequential( Conv2d(1, 32, kernel_size=3, stride=1, padding=1), Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True), Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True), Conv2d(32, 64, kernel_size=3, stride=(3, 1), padding=1), Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True), Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True), Conv2d(64, 128, kernel_size=3, stride=3, padding=1), Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True), Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True), Conv2d(128, 256, kernel_size=3, stride=(3, 2), padding=1), Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True), Conv2d(256, 512, kernel_size=3, stride=1, padding=0), Conv2d(512, 512, kernel_size=1, stride=1, padding=0), ) for child in self.audio_encoder_liujin.children(): for param in child.parameters(): param.requires_grad = False # self.lstm = nn.LSTM(256*4,256,3,batch_first = True) self.lstm_liujin = nn.LSTM(256 * 5, 256, 3, batch_first=True) # self.lstm_fc = nn.Sequential( # nn.Linear(256,16), # ) self.decon = nn.Sequential( nn.ConvTranspose2d(256, 256, kernel_size=6, stride=2, padding=1, bias=True),#4,4 nn.BatchNorm2d(256), nn.ReLU(True), nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1, bias=True),#8,8 nn.BatchNorm2d(128), nn.ReLU(True), nn.ConvTranspose2d(128, 128, kernel_size=4, stride=2, padding=1, bias=True), #16,16 nn.BatchNorm2d(128), nn.ReLU(True), nn.ConvTranspose2d(128, 128, kernel_size=4, stride=2, padding=1, bias=True),#32,32 nn.BatchNorm2d(128), nn.ReLU(True), nn.ConvTranspose2d(128, 32+3, kernel_size=4, stride=2, padding=1, bias=True),#64,64 # nn.ConvTranspose2d(128, 32*4, kernel_size=2, stride=2, padding=3, bias=True),#64,64 ) self.generator = Generator(64,256,8) def forward(self, example_image, audio, pose, jaco_net): hidden = ( torch.autograd.Variable(torch.zeros(3, audio.size(0), 256).cuda()), torch.autograd.Variable(torch.zeros(3, audio.size(0), 256).cuda())) outs = example_image for down_block in self.down_blocks: outs = down_block(outs) image_feature = outs image_feature = image_feature.view(image_feature.shape[0], -1) lstm_input = [] for step_t in range(audio.size(1)): current_audio = audio[ : ,step_t , :, :].unsqueeze(1) # current_feature = self.audio_eocder(current_audio) # current_feature = current_feature.view(current_feature.size(0), -1) # current_feature = self.audio_eocder_fc(current_feature) current_feature = self.audio_encoder_liujin(current_audio) current_feature = current_feature.reshape(current_audio.shape[0], -1) pose_f = self.pose_encoder(pose[:,step_t]) features = torch.cat([image_feature, current_feature, pose_f], 1) lstm_input.append(features) lstm_input = torch.stack(lstm_input, dim = 1) # lstm_out, _ = self.lstm(lstm_input, hidden) lstm_out, _ = self.lstm_liujin(lstm_input, hidden) fc_out = [] deco_out = [] for step_t in range(audio.size(1)): fc_in = lstm_out[:,step_t,:] # fc_out.append(self.lstm_fc(fc_in)) if jaco_net == 'cnn': fc_feature = torch.unsqueeze(fc_in,2) fc_feature = torch.unsqueeze(fc_feature,3) deco_out.append(self.decon(fc_feature)) elif jaco_net == 'gan': result,_ = self.generator([fc_in]) deco_out.append(result) else: raise Exception("jaco_net type wrong") return torch.stack(deco_out,dim=1) class Classify(nn.Module): def __init__(self): super(Classify, self).__init__() self.last_fc = nn.Linear(512,8) def forward(self, feature): # mfcc= torch.unsqueeze(mfcc, 1) x = self.last_fc(feature) return x class TF_net(nn.Module): def __init__(self): super(TF_net, self).__init__() down_blocks = [] for i in range(8): down_blocks.append(DownBlock2d(3 if i == 0 else 2 * (2 ** i), 2 * (2 ** (i + 1)), kernel_size=3, padding=1)) self.down_blocks = nn.ModuleList(down_blocks) # self.lmark_encoder = nn.Sequential( # nn.Linear(16,256), # nn.ReLU(True), # nn.Linear(256,512), # nn.ReLU(True), # ) self.pose_encoder = nn.Sequential( nn.Linear(6,128), nn.ReLU(True), nn.Linear(128,256), nn.ReLU(True), ) self.audio_eocder = nn.Sequential( conv2d(1,64,3,1,1), conv2d(64,128,3,1,1), nn.MaxPool2d(3, stride=(1,2)), conv2d(128,256,3,1,1), conv2d(256,256,3,1,1), conv2d(256,512,3,1,1), nn.MaxPool2d(3, stride=(2,2)) ) self.audio_eocder_fc = nn.Sequential( nn.Linear(1024 *12,2048), nn.ReLU(True), nn.Linear(2048,256), nn.ReLU(True), ) self.lstm = nn.LSTM(256*4,256,3,batch_first = True) self.lstm_two = nn.LSTM(256*6,256,3,batch_first = True) # self.lstm_fc = nn.Sequential( # nn.Linear(256,16), # ) self.decon = nn.Sequential( nn.ConvTranspose2d(256, 256, kernel_size=6, stride=2, padding=1, bias=True),#4,4 nn.BatchNorm2d(256), nn.ReLU(True), nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1, bias=True),#8,8 nn.BatchNorm2d(128), nn.ReLU(True), nn.ConvTranspose2d(128, 128, kernel_size=4, stride=2, padding=1, bias=True), #16,16 nn.BatchNorm2d(128), nn.ReLU(True), nn.ConvTranspose2d(128, 128, kernel_size=4, stride=2, padding=1, bias=True),#32,32 nn.BatchNorm2d(128), nn.ReLU(True), nn.ConvTranspose2d(128, 32+3, kernel_size=4, stride=2, padding=1, bias=True),#64,64 # nn.ConvTranspose2d(128, 32*4, kernel_size=2, stride=2, padding=3, bias=True),#64,64 ) self.generator = Generator(64,256,8) self.instance_norm = InstanceNorm() self.style_mod = ApplyStyle(512, 1024, use_wscale=True) self.style_mod1 = ApplyStyle(512, 35, use_wscale=True) def adain_forward(self, example_image, audio, pose, jaco_net, emo_features): hidden = ( torch.autograd.Variable(torch.zeros(3, audio.size(0), 256).cuda()), torch.autograd.Variable(torch.zeros(3, audio.size(0), 256).cuda())) outs = example_image for down_block in self.down_blocks: outs = down_block(outs) image_feature = outs image_feature = image_feature.view(image_feature.shape[0], -1) lstm_input = [] for step_t in range(audio.size(1)): current_audio = audio[ : ,step_t , :, :].unsqueeze(1) current_feature = self.audio_eocder(current_audio) current_feature = current_feature.view(current_feature.size(0), -1) current_feature = self.audio_eocder_fc(current_feature) #256 pose_f = self.pose_encoder(pose[:,step_t]) #256 features = torch.cat([image_feature, current_feature, pose_f], 1) features = torch.unsqueeze(torch.unsqueeze(features,-1),-1) features = self.instance_norm(features) x = self.style_mod(features, emo_features[step_t]) # t = adain(torch.unsqueeze(torch.unsqueeze(features,-1),-1), torch.unsqueeze(torch.unsqueeze(emo_features[step_t],1),2)) lstm_input.append(torch.squeeze(torch.squeeze(x,-1),-1)) lstm_input = torch.stack(lstm_input, dim = 1) lstm_out, _ = self.lstm(lstm_input, hidden) # fc_out = [] deco_out = [] for step_t in range(audio.size(1)): fc_in = lstm_out[:,step_t,:] # fc_out.append(self.lstm_fc(fc_in)) if jaco_net == 'cnn': fc_feature = torch.unsqueeze(fc_in,2) fc_feature = torch.unsqueeze(fc_feature,3) deco_out.append(self.decon(fc_feature)) elif jaco_net == 'gan': result,_ = self.generator([fc_in]) deco_out.append(result) else: raise Exception("jaco_net type wrong") return torch.stack(deco_out,dim=1) def adain_feature2(self, example_image, audio, pose, jaco_net, emo_features): hidden = ( torch.autograd.Variable(torch.zeros(3, audio.size(0), 256).cuda()), torch.autograd.Variable(torch.zeros(3, audio.size(0), 256).cuda())) outs = example_image for down_block in self.down_blocks: outs = down_block(outs) image_feature = outs image_feature = image_feature.view(image_feature.shape[0], -1) lstm_input = [] for step_t in range(audio.size(1)): current_audio = audio[ : ,step_t , :, :].unsqueeze(1) current_feature = self.audio_eocder(current_audio) current_feature = current_feature.view(current_feature.size(0), -1) current_feature = self.audio_eocder_fc(current_feature) #256 pose_f = self.pose_encoder(pose[:,step_t]) #256 features = torch.cat([image_feature, current_feature, pose_f], 1) lstm_input.append(features) lstm_input = torch.stack(lstm_input, dim = 1) lstm_out, _ = self.lstm(lstm_input, hidden) # fc_out = [] deco_out = [] for step_t in range(audio.size(1)): fc_in = lstm_out[:,step_t,:] # fc_out.append(self.lstm_fc(fc_in)) if jaco_net == 'cnn': fc_feature = torch.unsqueeze(fc_in,2) fc_feature = torch.unsqueeze(fc_feature,3) fc_feature = self.decon(fc_feature) fc_feature = self.instance_norm(fc_feature) t = self.style_mod1(fc_feature, emo_features[step_t]) # emo_feature = torch.unsqueeze(torch.unsqueeze(emo_features[step_t],-1),-1) # emo_feature = emo_feature.repeat(1,fc_feature.shape[1],1,1) # t = adain(fc_feature, emo_feature) deco_out.append(t) elif jaco_net == 'gan': result,_ = self.generator([fc_in]) deco_out.append(result) else: raise Exception("jaco_net type wrong") return torch.stack(deco_out,dim=1) def forward(self, example_image, audio, pose, jaco_net, emo_features): hidden = ( torch.autograd.Variable(torch.zeros(3, audio.size(0), 256).cuda()), torch.autograd.Variable(torch.zeros(3, audio.size(0), 256).cuda())) outs = example_image for down_block in self.down_blocks: outs = down_block(outs) image_feature = outs image_feature = image_feature.view(image_feature.shape[0], -1) lstm_input = [] for step_t in range(audio.size(1)): current_audio = audio[ : ,step_t , :, :].unsqueeze(1) current_feature = self.audio_eocder(current_audio) current_feature = current_feature.view(current_feature.size(0), -1) current_feature = self.audio_eocder_fc(current_feature) #256 pose_f = self.pose_encoder(pose[:,step_t]) #256 features = torch.cat([image_feature, current_feature, pose_f, emo_features[step_t]], 1) lstm_input.append(features) lstm_input = torch.stack(lstm_input, dim = 1) lstm_out, _ = self.lstm_two(lstm_input, hidden) fc_out = [] deco_out = [] for step_t in range(audio.size(1)): fc_in = lstm_out[:,step_t,:] # fc_out.append(self.lstm_fc(fc_in)) if jaco_net == 'cnn': fc_feature = torch.unsqueeze(fc_in,2) fc_feature = torch.unsqueeze(fc_feature,3) deco_out.append(self.decon(fc_feature)) elif jaco_net == 'gan': result,_ = self.generator([fc_in]) deco_out.append(result) else: raise Exception("jaco_net type wrong") return torch.stack(deco_out,dim=1) class AT_net2_liujin(nn.Module): def __init__(self): super(AT_net2_liujin, self).__init__() down_blocks = [] for i in range(8): down_blocks.append(DownBlock2d(3 if i == 0 else 2 * (2 ** i), 2 * (2 ** (i + 1)), kernel_size=3, padding=1)) self.down_blocks = nn.ModuleList(down_blocks) # self.lmark_encoder = nn.Sequential( # nn.Linear(16,256), # nn.ReLU(True), # nn.Linear(256,512), # nn.ReLU(True), # ) self.pose_encoder = nn.Sequential( nn.Linear(6,128), nn.ReLU(True), nn.Linear(128,256), nn.ReLU(True), ) self.audio_eocder = nn.Sequential( conv2d(1,64,3,1,1), conv2d(64,128,3,1,1), nn.MaxPool2d(3, stride=(1,2)), conv2d(128,256,3,1,1), conv2d(256,256,3,1,1), conv2d(256,512,3,1,1), nn.MaxPool2d(3, stride=(2,2)) ) self.audio_eocder_fc = nn.Sequential( nn.Linear(1024 *12,2048), nn.ReLU(True), nn.Linear(2048,256), nn.ReLU(True), ) # self.lstm = nn.LSTM(256*4,256,3,batch_first = True) self.lstm_liujin = nn.LSTM(256 * 5, 256, 3, batch_first=True) # self.lstm_fc = nn.Sequential( # nn.Linear(256,16), # ) self.decon = nn.Sequential( nn.ConvTranspose2d(256, 256, kernel_size=6, stride=2, padding=1, bias=True),#4,4 nn.BatchNorm2d(256), nn.ReLU(True), nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1, bias=True),#8,8 nn.BatchNorm2d(128), nn.ReLU(True), nn.ConvTranspose2d(128, 128, kernel_size=4, stride=2, padding=1, bias=True), #16,16 nn.BatchNorm2d(128), nn.ReLU(True), nn.ConvTranspose2d(128, 128, kernel_size=4, stride=2, padding=1, bias=True),#32,32 nn.BatchNorm2d(128), nn.ReLU(True), nn.ConvTranspose2d(128, 32+3, kernel_size=4, stride=2, padding=1, bias=True),#64,64 # nn.ConvTranspose2d(128, 32*4, kernel_size=2, stride=2, padding=3, bias=True),#64,64 ) self.generator = Generator(64,256,8) def forward(self, example_image, audio, pose, jaco_net, weight): hidden = ( torch.autograd.Variable(torch.zeros(3, audio.size(0), 256).cuda()), torch.autograd.Variable(torch.zeros(3, audio.size(0), 256).cuda())) outs = example_image for down_block in self.down_blocks: outs = down_block(outs) image_feature = outs image_feature = image_feature.view(image_feature.shape[0], -1) #1, 512, 1, 1 lstm_input = [] for step_t in range(audio.size(1)): current_audio = audio[ : ,step_t , :, :].unsqueeze(1) current_feature = self.audio_eocder(current_audio) current_feature = current_feature.view(current_feature.size(0), -1) current_feature = self.audio_eocder_fc(current_feature)*weight pose_f = self.pose_encoder(pose[:,step_t]) features = torch.cat([image_feature, current_feature, pose_f], 1) lstm_input.append(features) lstm_input = torch.stack(lstm_input, dim = 1) lstm_out, _ = self.lstm(lstm_input, hidden) #1,N,256 fc_out = [] deco_out = [] for step_t in range(audio.size(1)): fc_in = lstm_out[:,step_t,:] # fc_out.append(self.lstm_fc(fc_in)) if jaco_net == 'cnn': fc_feature = torch.unsqueeze(fc_in,2) fc_feature = torch.unsqueeze(fc_feature,3) deco_out.append(self.decon(fc_feature)) elif jaco_net == 'gan': result,_ = self.generator([fc_in]) deco_out.append(result) else: raise Exception("jaco_net type wrong") return torch.stack(deco_out,dim=1) class Ct_encoder(nn.Module): def __init__(self): super(Ct_encoder, self).__init__() self.audio_eocder = nn.Sequential( conv2d(1,64,3,1,1), conv2d(64,128,3,1,1), nn.MaxPool2d(3, stride=(1,2)), conv2d(128,256,3,1,1), conv2d(256,256,3,1,1), conv2d(256,512,3,1,1), nn.MaxPool2d(3, stride=(2,2)) ) self.audio_eocder_fc = nn.Sequential( nn.Linear(1024 *12,2048), nn.ReLU(True), nn.Linear(2048,256), nn.ReLU(True), ) def forward(self, audio): feature = self.audio_eocder(audio) feature = feature.view(feature.size(0),-1) x = self.audio_eocder_fc(feature) return x class EmotionNet(nn.Module): def __init__(self): super(EmotionNet, self).__init__() self.emotion_eocder = nn.Sequential( conv2d(1,64,3,1,1), nn.MaxPool2d((1,3), stride=(1,2)), #[1, 64, 12, 12] conv2d(64,128,3,1,1), conv2d(128,256,3,1,1), nn.MaxPool2d((12,1), stride=(12,1)), #[1, 256, 1, 12] conv2d(256,512,3,1,1), nn.MaxPool2d((1,2), stride=(1,2)) #[1, 512, 1, 6] ) self.emotion_eocder_fc = nn.Sequential( nn.Linear(512 *6,2048), nn.ReLU(True), nn.Linear(2048,128), nn.ReLU(True), ) self.last_fc = nn.Linear(128,8) self.re_id = nn.Sequential( conv2d(512,1024,3,1,1), nn.MaxPool2d((1,2), stride=(1,2)), #[1, 1024, 1, 3] conv2d(1024,1024,3,1,1), conv2d(1024,2048,3,1,1), nn.MaxPool2d((1,2), stride=(1,2)) #[1, 2048, 1, 1] ) self.re_id_fc = nn.Sequential( nn.Linear(2048,512), nn.ReLU(True), nn.Linear(512,128), nn.ReLU(True), ) def forward(self, mfcc): # mfcc= torch.unsqueeze(mfcc, 1) mfcc=torch.transpose(mfcc,2,3) feature = self.emotion_eocder(mfcc) # id_feature = feature.detach() feature = feature.view(feature.size(0),-1) x = self.emotion_eocder_fc(feature) # remove_feature = self.re_id(id_feature) # remove_feature = remove_feature.view(remove_feature.size(0),-1) # y = self.re_id_fc(remove_feature) return x class AF2F(nn.Module): def __init__(self): super(AF2F, self).__init__() self.decon = nn.Sequential( nn.ConvTranspose2d(384, 256, kernel_size=6, stride=2, padding=1, bias=True),#4,4 nn.BatchNorm2d(256), nn.ReLU(True), nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1, bias=True),#8,8 nn.BatchNorm2d(128), nn.ReLU(True), nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1, bias=True), #16,16 nn.BatchNorm2d(64), nn.ReLU(True), nn.ConvTranspose2d(64, 64, kernel_size=4, stride=2, padding=1, bias=True),#32,32 nn.BatchNorm2d(64), nn.ReLU(True), nn.ConvTranspose2d(64, 32+3, kernel_size=4, stride=2, padding=1, bias=True),#64,64 ) def forward(self, content,emotion): features = torch.cat([content, emotion], 1) #connect tensors inputs and dimension features = torch.unsqueeze(features,2) features = torch.unsqueeze(features,3) x = self.decon(features) return x class AF2F_s(nn.Module): def __init__(self): super(AF2F_s, self).__init__() self.decon = nn.Sequential( nn.ConvTranspose2d(256, 256, kernel_size=6, stride=2, padding=1, bias=True),#4,4 nn.BatchNorm2d(256), nn.ReLU(True), nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1, bias=True),#8,8 nn.BatchNorm2d(128), nn.ReLU(True), nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1, bias=True), #16,16 nn.BatchNorm2d(64), nn.ReLU(True), nn.ConvTranspose2d(64, 64, kernel_size=4, stride=2, padding=1, bias=True),#32,32 nn.BatchNorm2d(64), nn.ReLU(True), nn.ConvTranspose2d(64, 32+3, kernel_size=4, stride=2, padding=1, bias=True),#64,64 nn.ReLU(), ) def forward(self, content): # features = torch.cat([content, emotion], 1) #connect tensors inputs and dimension features = torch.unsqueeze(content,2) features = torch.unsqueeze(features,3) x = self.decon(features) return x class A2I(nn.Module): def __init__(self): super(A2I, self).__init__() self.audio_eocder = nn.Sequential( conv2d(1,64,3,1,1), conv2d(64,128,3,1,1), nn.MaxPool2d((1,5), stride=(1,2)), conv2d(128,256,3,1,1), conv2d(256,256,3,1,1), nn.MaxPool2d((5,5), stride=(2,2)) ) self.decon = nn.Sequential( nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1, bias=True),#8,8 nn.BatchNorm2d(128), nn.ReLU(True), nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1, bias=True), #16,16 nn.BatchNorm2d(64), nn.ReLU(True), nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1, bias=True),#32,32 nn.BatchNorm2d(32), nn.ReLU(True), nn.ConvTranspose2d(32, 2, kernel_size=4, stride=2, padding=1, bias=True),#64,64 nn.ReLU(), ) def forward(self, mfcc): mfcc= torch.unsqueeze(mfcc, 1) mfcc=torch.transpose(mfcc,2,3) feature = self.audio_eocder(mfcc) # id_feature = feature.detach() x = self.decon(feature) return x def kp2gaussian(kp, spatial_size, kp_variance): """ Transform a keypoint into gaussian like representation """ mean = kp['value'] #[4,10,2] coordinate_grid = make_coordinate_grid(spatial_size, mean.type()) #[h,w,2] number_of_leading_dimensions = len(mean.shape) - 1 shape = (1,) * number_of_leading_dimensions + coordinate_grid.shape #5 coordinate_grid = coordinate_grid.view(*shape) #[1,1,h,w,2] repeats = mean.shape[:number_of_leading_dimensions] + (1, 1, 1) coordinate_grid = coordinate_grid.repeat(*repeats) #[4,10,h,w,2] # Preprocess kp shape shape = mean.shape[:number_of_leading_dimensions] + (1, 1, 2) mean = mean.view(*shape) #[4,10,1,1,2] 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, 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 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 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 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 SameBlock2d(nn.Module): """ Simple block, preserve spatial resolution. """ def __init__(self, in_features, out_features, groups=1, kernel_size=3, padding=1): 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) def forward(self, x): out = self.conv(x) out = self.norm(out) out = F.relu(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(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)) 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(UpBlock2d(in_filters, out_filters, kernel_size=3, padding=1)) self.up_blocks = nn.ModuleList(up_blocks) self.out_filters = block_expansion + in_features 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) 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 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 sigma = 1.5 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 def sigmoid(x): return 1 / (1 + math.exp(-x)) def norm_angle(angle): norm_angle = sigmoid(10 * (abs(angle) / 0.7853975 - 1)) return norm_angle def conv3x3(in_planes, out_planes, stride=1): "3x3 convolution with padding" return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False) class BasicBlock(nn.Module): expansion = 1 def __init__(self, inplanes, planes, stride=1, downsample=None): super(BasicBlock, self).__init__() self.conv1 = conv3x3(inplanes, planes, stride) self.bn1 = nn.BatchNorm2d(planes) self.relu = nn.ReLU() self.conv2 = conv3x3(planes, planes) self.bn2 = nn.BatchNorm2d(planes) self.downsample = downsample self.stride = stride def forward(self, x): residual = x out = self.conv1(x) out = self.bn1(out) out = self.relu(out) out = self.conv2(out) out = self.bn2(out) if self.downsample is not None: residual = self.downsample(x) out += residual out = self.relu(out) return out class Bottleneck(nn.Module): expansion = 4 def __init__(self, inplanes, planes, stride=1, downsample=None): super(Bottleneck, self).__init__() self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) self.bn1 = nn.BatchNorm2d(planes) self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False) self.bn2 = nn.BatchNorm2d(planes) self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) self.bn3 = nn.BatchNorm2d(planes * 4) self.relu = nn.ReLU() self.downsample = downsample self.stride = stride def forward(self, x): residual = x out = self.conv1(x) out = self.bn1(out) out = self.relu(out) out = self.conv2(out) out = self.bn2(out) out = self.relu(out) out = self.conv3(out) out = self.bn3(out) if self.downsample is not None: residual = self.downsample(x) out = out + residual out = self.relu(out) return out class EmDetector(nn.Module): """ Detecting a keypoints. Return keypoint position and jacobian near each keypoint. """ def __init__(self, block_expansion, num_channels, max_features, num_blocks, scale_factor=1, num_classes=8): super(EmDetector, self).__init__() self.inplanes = 64 self.predictor = Hourglass(block_expansion, in_features=num_channels, max_features=max_features, num_blocks=num_blocks) self.scale_factor = scale_factor if self.scale_factor != 1: self.down = AntiAliasInterpolation2d(num_channels, self.scale_factor) self.conv1 = nn.Conv2d(self.predictor.out_filters, 64, kernel_size=3, stride=1, padding=1, bias=False) self.bn1 = nn.BatchNorm2d(64) self.relu = nn.ReLU() self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) layers = [2,2,2,2] self.layer1 = self._make_layer(BasicBlock, 64, layers[0]) self.layer2 = self._make_layer(BasicBlock, 128, layers[1], stride=2) self.layer3 = self._make_layer(BasicBlock, 256, layers[2], stride=2) self.layer4 = self._make_layer(BasicBlock, 512, layers[3], stride=2) self.avgpool = nn.AdaptiveAvgPool2d(1) self.fc = nn.Linear(512 * BasicBlock.expansion, num_classes) self.classify = Classify() def _make_layer(self, block, planes, blocks, stride=1): downsample = None if stride != 1 or self.inplanes != planes * block.expansion: downsample = nn.Sequential( nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, bias=False), nn.BatchNorm2d(planes * block.expansion), ) layers = [] layers.append(block(self.inplanes, planes, stride, downsample)) self.inplanes = planes * block.expansion for i in range(1, blocks): layers.append(block(self.inplanes, planes)) return nn.Sequential(*layers) def adain_feature(self, x): #torch.Size([4, 3, H, W]) if self.scale_factor != 1: x = self.down(x) # 0.25 [4, 3, H/4, W/4] feature_map = self.predictor(x) #[4,3+32,H/4, W/4] # out = self.fc(out) return feature_map def forward(self, x): #torch.Size([4, 3, H, W]) if self.scale_factor != 1: x = self.down(x) # 0.25 [4, 3, H/4, W/4] feature_map = self.predictor(x) #[4,3+32,H/4, W/4] f = self.conv1(feature_map) #[16,64,64,64] f = self.bn1(f) #torch.Size([16, 64, 64, 64]) f = self.relu(f) f = self.maxpool(f) #[16, 64, 32, 32] f = self.layer1(f) #[16, 64, 32, 32] f = self.layer2(f) #[16, 128, 16, 16]) f = self.layer3(f) #[16, 256, 8, 8] f = self.layer4(f) #[16, 512, 4, 4] f = self.avgpool(f) #[16, 512, 1, 1] out = f.squeeze(3).squeeze(2) fake = self.classify(out) # out = self.fc(out) return out, fake class Emotion_k(nn.Module): """ Detecting a keypoints. Return keypoint position and jacobian near each keypoint. """ def __init__(self, block_expansion, num_channels, max_features, num_blocks, scale_factor=1, num_classes=8): super(Emotion_k, self).__init__() self.inplanes = 64 self.predictor = Hourglass(block_expansion, in_features=num_channels, max_features=max_features, num_blocks=num_blocks) self.scale_factor = scale_factor if self.scale_factor != 1: self.down = AntiAliasInterpolation2d(num_channels, self.scale_factor) self.conv1 = nn.Conv2d(self.predictor.out_filters, 64, kernel_size=3, stride=1, padding=1, bias=False) self.bn1 = nn.BatchNorm2d(64) self.relu = nn.ReLU() self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) layers = [2,2,2,2] self.layer1 = self._make_layer(BasicBlock, 64, layers[0]) self.layer2 = self._make_layer(BasicBlock, 128, layers[1], stride=2) self.layer3 = self._make_layer(BasicBlock, 256, layers[2], stride=2) self.layer4 = self._make_layer(BasicBlock, 512, layers[3], stride=2) self.avgpool = nn.AdaptiveAvgPool2d(1) self.fc = nn.Linear(512 * BasicBlock.expansion, num_classes) self.embed_fn, self.input_ch = get_embedder(10, 0) self.fc_p = nn.Sequential( nn.Linear(10 * 126,1024), nn.ReLU(True), nn.Linear(1024,512), nn.ReLU(True), ) self.fc_n = nn.Sequential( nn.Linear(10 * 6,128), nn.ReLU(True), nn.Linear(128,512), nn.ReLU(True), ) self.fc_all = nn.Sequential( nn.Linear(1024,512), nn.ReLU(True), nn.Linear(512,256), nn.ReLU(True), nn.Linear(256,64), nn.ReLU(True), ) # self.fc_single = nn.Sequential( # nn.Linear(512,256), # nn.ReLU(True), # nn.Linear(256,64), # nn.ReLU(True), # ) self.final = nn.Sequential( nn.Conv1d(1,2,4,2,1), nn.MaxPool1d(2,stride=2), nn.ReLU(True), nn.Conv1d(2,4,4,2,1), nn.ReLU(True), nn.Conv1d(4,4,3), ) self.final_4 = nn.Sequential( nn.Conv1d(4,4,3,1,1), nn.MaxPool1d(2,stride=2), nn.ReLU(True), nn.Conv1d(4,4,3,1) ) self.final_10 = nn.Sequential( nn.Conv1d(4,8,3,1,1), #[B,8,16] nn.MaxPool1d(2,stride=2), #[B,8,8] nn.ReLU(True), nn.Conv1d(8,10,3,1), #[B,10,6] ) self.classify = Classify() def _make_layer(self, block, planes, blocks, stride=1): downsample = None if stride != 1 or self.inplanes != planes * block.expansion: downsample = nn.Sequential( nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, bias=False), nn.BatchNorm2d(planes * block.expansion), ) layers = [] layers.append(block(self.inplanes, planes, stride, downsample)) self.inplanes = planes * block.expansion for i in range(1, blocks): layers.append(block(self.inplanes, planes)) return nn.Sequential(*layers) def linear_10(self, x, value, jacobian): #torch.Size([4, 3, H, W]) if self.scale_factor != 1: x = self.down(x) # 0.25 [4, 3, H/4, W/4] feature_map = self.predictor(x) #[4,3+32,H/4, W/4] f = self.conv1(feature_map) #[16,64,64,64] f = self.bn1(f) #torch.Size([16, 64, 64, 64]) f = self.relu(f) f = self.maxpool(f) #[16, 64, 32, 32] f = self.layer1(f) #[16, 64, 32, 32] f = self.layer2(f) #[16, 128, 16, 16]) f = self.layer3(f) #[16, 256, 8, 8] f = self.layer4(f) #[16, 512, 4, 4] f = self.avgpool(f) #[16, 512, 1, 1] out = f.squeeze(3).squeeze(2) fake = self.classify(out) jacobian = jacobian.reshape(jacobian.shape[0],jacobian.shape[1],4) neu_input = torch.cat((value,jacobian),2) posi_input = self.embed_fn(neu_input) posi_input =posi_input.reshape(posi_input.shape[0],-1) ner_feature = self.fc_p(posi_input) all_fc = self.fc_all(torch.cat((out,ner_feature),1)).reshape(-1,4,16) result = self.final_10(all_fc) e_value = result[:,:,:2] e_jacobian = result[:,:,2:].reshape(result.shape[0],10,2,2) kp = {'value': e_value,'jacobian': e_jacobian} return kp, fake def linear_4(self, x, value, jacobian): #torch.Size([4, 3, H, W]) if self.scale_factor != 1: x = self.down(x) # 0.25 [4, 3, H/4, W/4] feature_map = self.predictor(x) #[4,3+32,H/4, W/4] f = self.conv1(feature_map) #[16,64,64,64] f = self.bn1(f) #torch.Size([16, 64, 64, 64]) f = self.relu(f) f = self.maxpool(f) #[16, 64, 32, 32] f = self.layer1(f) #[16, 64, 32, 32] f = self.layer2(f) #[16, 128, 16, 16]) f = self.layer3(f) #[16, 256, 8, 8] f = self.layer4(f) #[16, 512, 4, 4] f = self.avgpool(f) #[16, 512, 1, 1] out = f.squeeze(3).squeeze(2) fake = self.classify(out) # jacobian = jacobian.reshape(jacobian.shape[0],jacobian.shape[1],4) # neu_input = torch.cat((value,jacobian),2) # posi_input = self.embed_fn(neu_input) # posi_input =posi_input.reshape(posi_input.shape[0],-1) # ner_feature = self.fc_p(posi_input) # all_fc = self.fc_all(torch.cat((out,ner_feature),1)).reshape(-1,4,16) all_fc = torch.unsqueeze(self.fc_single(out),1) result = self.final(all_fc) e_value = result[:,:,:2] e_jacobian = result[:,:,2:].reshape(result.shape[0],4,2,2) kp = {'value': e_value,'jacobian': e_jacobian} # out = self.fc(out) return kp, fake def linear_np_10(self, x, value, jacobian): #torch.Size([4, 3, H, W]) if self.scale_factor != 1: x = self.down(x) # 0.25 [4, 3, H/4, W/4] feature_map = self.predictor(x) #[4,3+32,H/4, W/4] f = self.conv1(feature_map) #[16,64,64,64] f = self.bn1(f) #torch.Size([16, 64, 64, 64]) f = self.relu(f) f = self.maxpool(f) #[16, 64, 32, 32] f = self.layer1(f) #[16, 64, 32, 32] f = self.layer2(f) #[16, 128, 16, 16]) f = self.layer3(f) #[16, 256, 8, 8] f = self.layer4(f) #[16, 512, 4, 4] f = self.avgpool(f) #[16, 512, 1, 1] out = f.squeeze(3).squeeze(2) fake = self.classify(out) jacobian = jacobian.reshape(jacobian.shape[0],jacobian.shape[1],4) neu_input = torch.cat((value,jacobian),2) posi_input =neu_input.reshape(neu_input.shape[0],-1) ner_feature = self.fc_n(posi_input) all_fc = self.fc_all(torch.cat((out,ner_feature),1)).reshape(-1,4,16) result = self.final_10(all_fc) e_value = result[:,:,:2] e_jacobian = result[:,:,2:].reshape(result.shape[0],10,2,2) kp = {'value': e_value,'jacobian': e_jacobian} # out = self.fc(out) return kp, fake def linear_np_4(self, x, value, jacobian): #torch.Size([4, 3, H, W]) if self.scale_factor != 1: x = self.down(x) # 0.25 [4, 3, H/4, W/4] feature_map = self.predictor(x) #[4,3+32,H/4, W/4] f = self.conv1(feature_map) #[16,64,64,64] f = self.bn1(f) #torch.Size([16, 64, 64, 64]) f = self.relu(f) f = self.maxpool(f) #[16, 64, 32, 32] f = self.layer1(f) #[16, 64, 32, 32] f = self.layer2(f) #[16, 128, 16, 16]) f = self.layer3(f) #[16, 256, 8, 8] f = self.layer4(f) #[16, 512, 4, 4] f = self.avgpool(f) #[16, 512, 1, 1] out = f.squeeze(3).squeeze(2) fake = self.classify(out) jacobian = jacobian.reshape(jacobian.shape[0],jacobian.shape[1],4) neu_input = torch.cat((value,jacobian),2) posi_input =neu_input.reshape(neu_input.shape[0],-1) ner_feature = self.fc_n(posi_input) all_fc = torch.unsqueeze(self.fc_all(torch.cat((out,ner_feature),1)),1) result = self.final(all_fc) e_value = result[:,:,:2] e_jacobian = result[:,:,2:].reshape(result.shape[0],4,2,2) kp = {'value': e_value,'jacobian': e_jacobian} # out = self.fc(out) return kp, fake def emotion_feature(self, feature, value, jacobian): #torch.Size([4, 3, H, W]) out = feature fake = self.classify(out) jacobian = jacobian.reshape(jacobian.shape[0],jacobian.shape[1],4) neu_input = torch.cat((value,jacobian),2) posi_input = self.embed_fn(neu_input) posi_input =posi_input.reshape(posi_input.shape[0],-1) ner_feature = self.fc_p(posi_input) all_fc = torch.unsqueeze(self.fc_all(torch.cat((out,ner_feature),1)),1) result = self.final(all_fc) e_value = result[:,:,:2] e_jacobian = result[:,:,2:].reshape(result.shape[0],4,2,2) kp = {'value': e_value,'jacobian': e_jacobian} # out = self.fc(out) return kp, fake def feature(self, x): #torch.Size([4, 3, H, W]) if self.scale_factor != 1: x = self.down(x) # 0.25 [4, 3, H/4, W/4] feature_map = self.predictor(x) #[4,3+32,H/4, W/4] f = self.conv1(feature_map) #[16,64,64,64] f = self.bn1(f) #torch.Size([16, 64, 64, 64]) f = self.relu(f) f = self.maxpool(f) #[16, 64, 32, 32] f = self.layer1(f) #[16, 64, 32, 32] f = self.layer2(f) #[16, 128, 16, 16]) f = self.layer3(f) #[16, 256, 8, 8] f = self.layer4(f) #[16, 512, 4, 4] f = self.avgpool(f) #[16, 512, 1, 1] out = f.squeeze(3).squeeze(2) # out = self.fc(out) return out def forward(self, x, value, jacobian): #torch.Size([4, 3, H, W]) if self.scale_factor != 1: x = self.down(x) # 0.25 [4, 3, H/4, W/4] feature_map = self.predictor(x) #[4,3+32,H/4, W/4] f = self.conv1(feature_map) #[16,64,64,64] f = self.bn1(f) #torch.Size([16, 64, 64, 64]) f = self.relu(f) f = self.maxpool(f) #[16, 64, 32, 32] f = self.layer1(f) #[16, 64, 32, 32] f = self.layer2(f) #[16, 128, 16, 16]) f = self.layer3(f) #[16, 256, 8, 8] f = self.layer4(f) #[16, 512, 4, 4] f = self.avgpool(f) #[16, 512, 1, 1] out = f.squeeze(3).squeeze(2) fake = self.classify(out) jacobian = jacobian.reshape(jacobian.shape[0],jacobian.shape[1],4) neu_input = torch.cat((value,jacobian),2) posi_input = self.embed_fn(neu_input) posi_input =posi_input.reshape(posi_input.shape[0],-1) ner_feature = self.fc_p(posi_input) all_fc = torch.unsqueeze(self.fc_all(torch.cat((out,ner_feature),1)),1) result = self.final(all_fc) e_value = result[:,:,:2] e_jacobian = result[:,:,2:].reshape(result.shape[0],4,2,2) kp = {'value': e_value,'jacobian': e_jacobian} # out = self.fc(out) return kp, fake class Emotion_map(nn.Module): """ Detecting a keypoints. Return keypoint position and jacobian near each keypoint. """ def __init__(self, block_expansion, num_channels, max_features, num_blocks, scale_factor=1, num_classes=8): super(Emotion_map, self).__init__() self.inplanes = 64 self.predictor = Hourglass(block_expansion, in_features=num_channels, max_features=max_features, num_blocks=num_blocks) self.scale_factor = scale_factor if self.scale_factor != 1: self.down = AntiAliasInterpolation2d(num_channels, self.scale_factor) self.conv1 = nn.Conv2d(self.predictor.out_filters, 64, kernel_size=3, stride=1, padding=1, bias=False) self.bn1 = nn.BatchNorm2d(64) self.relu = nn.ReLU() self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) layers = [2,2,2,2] self.layer1 = self._make_layer(BasicBlock, 64, layers[0]) self.layer2 = self._make_layer(BasicBlock, 128, layers[1], stride=2) self.layer3 = self._make_layer(BasicBlock, 256, layers[2], stride=2) self.layer4 = self._make_layer(BasicBlock, 512, layers[3], stride=2) self.avgpool = nn.AdaptiveAvgPool2d(1) self.fc = nn.Linear(512 * BasicBlock.expansion, num_classes) self.embed_fn, self.input_ch = get_embedder(10, 0) self.fc_p = nn.Sequential( nn.Linear(10 * 126,1024), nn.ReLU(True), nn.Linear(1024,512), nn.ReLU(True), ) self.fc_all = nn.Sequential( nn.Linear(1024,2048), nn.ReLU(True) ) self.final = nn.Sequential( nn.ConvTranspose2d(128, 128, kernel_size=4, stride=2, padding=1, bias=True),#8,8 nn.BatchNorm2d(128), nn.ReLU(True), nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1, bias=True), #16,16 nn.BatchNorm2d(64), nn.ReLU(True), nn.ConvTranspose2d(64, 64, kernel_size=4, stride=2, padding=1, bias=True),#32,32 nn.BatchNorm2d(64), nn.ReLU(True), nn.ConvTranspose2d(64, 32+3, kernel_size=4, stride=2, padding=1, bias=True),#64,64 ) self.classify = Classify() self.kp = nn.Conv2d(in_channels=35, out_channels=10, kernel_size=(7, 7), padding=0) self.jacobian = nn.Conv2d(in_channels=35, out_channels=4 * 10, kernel_size=(7, 7), padding=0) self.jacobian.weight.data.zero_() self.jacobian.bias.data.copy_(torch.tensor([1, 0, 0, 1] * 10, dtype=torch.float)) self.temperature = 0.1 self.kp_4 = nn.Conv2d(in_channels=35, out_channels=4, kernel_size=(7, 7), padding=0) self.jacobian_4 = nn.Conv2d(in_channels=35, out_channels=4 * 4, kernel_size=(7, 7), padding=0) self.jacobian_4.weight.data.zero_() self.jacobian_4.bias.data.copy_(torch.tensor([1, 0, 0, 1] * 4, dtype=torch.float)) def _make_layer(self, block, planes, blocks, stride=1): downsample = None if stride != 1 or self.inplanes != planes * block.expansion: downsample = nn.Sequential( nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, bias=False), nn.BatchNorm2d(planes * block.expansion), ) layers = [] layers.append(block(self.inplanes, planes, stride, downsample)) self.inplanes = planes * block.expansion for i in range(1, blocks): layers.append(block(self.inplanes, planes)) return nn.Sequential(*layers) def gaussian2kp(self, heatmap): """ Extract the mean and from a heatmap """ shape = heatmap.shape heatmap = heatmap.unsqueeze(-1) #[4,10,58,58,1] grid = make_coordinate_grid(shape[2:], heatmap.type()).unsqueeze_(0).unsqueeze_(0) #[1,1,58,58,2] value = (heatmap * grid).sum(dim=(2, 3)) #[4,10,2] kp = {'value': value} return kp def map_4(self, x, value, jacobian): #torch.Size([4, 3, H, W]) if self.scale_factor != 1: x = self.down(x) # 0.25 [4, 3, H/4, W/4] feature_map = self.predictor(x) #[4,3+32,H/4, W/4] f = self.conv1(feature_map) #[16,64,64,64] f = self.bn1(f) #torch.Size([16, 64, 64, 64]) f = self.relu(f) f = self.maxpool(f) #[16, 64, 32, 32] f = self.layer1(f) #[16, 64, 32, 32] f = self.layer2(f) #[16, 128, 16, 16]) f = self.layer3(f) #[16, 256, 8, 8] f = self.layer4(f) #[16, 512, 4, 4] f = self.avgpool(f) #[16, 512, 1, 1] out = f.squeeze(3).squeeze(2) fake = self.classify(out) jacobian = jacobian.reshape(jacobian.shape[0],jacobian.shape[1],4) neu_input = torch.cat((value,jacobian),2) posi_input = self.embed_fn(neu_input) posi_input =posi_input.reshape(posi_input.shape[0],-1) ner_feature = self.fc_p(posi_input) all_fc = self.fc_all(torch.cat((out,ner_feature),1)).reshape(-1,128,4,4) feature_map = self.final(all_fc) prediction = self.kp_4(feature_map) #[4,10,H/4-6, W/4-6] final_shape = prediction.shape heatmap = prediction.view(final_shape[0], final_shape[1], -1) #[4, 10, 58*58] heatmap = F.softmax(heatmap / self.temperature, dim=2) heatmap = heatmap.view(*final_shape) #[4,10,58,58] out = self.gaussian2kp(heatmap) out['heatmap'] = heatmap if self.jacobian is not None: jacobian_map = self.jacobian_4(feature_map) ##[4,40,H/4-6, W/4-6] jacobian_map = jacobian_map.reshape(final_shape[0], 4, 4, final_shape[2], final_shape[3]) heatmap = heatmap.unsqueeze(2) jacobian = heatmap * jacobian_map #[4,10,4,H/4-6, W/4-6] jacobian = jacobian.view(final_shape[0], final_shape[1], 4, -1) jacobian = jacobian.sum(dim=-1) #[4,10,4] jacobian = jacobian.view(jacobian.shape[0], jacobian.shape[1], 2, 2) #[4,10,2,2] out['jacobian'] = jacobian return out, fake def forward(self, x, value, jacobian): #torch.Size([4, 3, H, W]) if self.scale_factor != 1: x = self.down(x) # 0.25 [4, 3, H/4, W/4] feature_map = self.predictor(x) #[4,3+32,H/4, W/4] f = self.conv1(feature_map) #[16,64,64,64] f = self.bn1(f) #torch.Size([16, 64, 64, 64]) f = self.relu(f) f = self.maxpool(f) #[16, 64, 32, 32] f = self.layer1(f) #[16, 64, 32, 32] f = self.layer2(f) #[16, 128, 16, 16]) f = self.layer3(f) #[16, 256, 8, 8] f = self.layer4(f) #[16, 512, 4, 4] f = self.avgpool(f) #[16, 512, 1, 1] out = f.squeeze(3).squeeze(2) fake = self.classify(out) jacobian = jacobian.reshape(jacobian.shape[0],jacobian.shape[1],4) neu_input = torch.cat((value,jacobian),2) posi_input = self.embed_fn(neu_input) posi_input =posi_input.reshape(posi_input.shape[0],-1) ner_feature = self.fc_p(posi_input) all_fc = self.fc_all(torch.cat((out,ner_feature),1)).reshape(-1,128,4,4) feature_map = self.final(all_fc) prediction = self.kp(feature_map) #[4,10,H/4-6, W/4-6] final_shape = prediction.shape heatmap = prediction.view(final_shape[0], final_shape[1], -1) #[4, 10, 58*58] heatmap = F.softmax(heatmap / self.temperature, dim=2) heatmap = heatmap.view(*final_shape) #[4,10,58,58] out = self.gaussian2kp(heatmap) out['heatmap'] = heatmap if self.jacobian is not None: jacobian_map = self.jacobian(feature_map) ##[4,40,H/4-6, W/4-6] jacobian_map = jacobian_map.reshape(final_shape[0], 10, 4, final_shape[2], final_shape[3]) heatmap = heatmap.unsqueeze(2) jacobian = heatmap * jacobian_map #[4,10,4,H/4-6, W/4-6] jacobian = jacobian.view(final_shape[0], final_shape[1], 4, -1) jacobian = jacobian.sum(dim=-1) #[4,10,4] jacobian = jacobian.view(jacobian.shape[0], jacobian.shape[1], 2, 2) #[4,10,2,2] out['jacobian'] = jacobian return out, fake def conv2d(channel_in, channel_out, ksize=3, stride=1, padding=1, activation=nn.ReLU, normalizer=nn.BatchNorm2d): layer = list() bias = True if not normalizer else False layer.append(nn.Conv2d(channel_in, channel_out, ksize, stride, padding, bias=bias)) _apply(layer, activation, normalizer, channel_out) # init.kaiming_normal(layer[0].weight) return nn.Sequential(*layer) def _apply(layer, activation, normalizer, channel_out=None): if normalizer: layer.append(normalizer(channel_out)) if activation: layer.append(activation()) return layer