Update test code

Author: yzhou359

main_end2end.py

@@ -35,18 +35,15 @@
 parser.add_argument('--jpg', type=str, default='{}.jpg'.format(default_head_name))
 parser.add_argument('--close_input_face_mouth', default=CLOSE_INPUT_FACE_MOUTH, action='store_true')
 
-
 parser.add_argument('--load_AUTOVC_name', type=str, default='examples/ckpt/ckpt_autovc.pth')
-parser.add_argument('--load_a2l_G_name', type=str, default='examples/ckpt/ckpt_a2l_db_e_875.pth')
-parser.add_argument('--load_a2l_C_name', type=str, default='examples/ckpt/ckpt_audio2landmark_c.pth')
+parser.add_argument('--load_a2l_G_name', type=str, default='examples/ckpt/ckpt_speaker_branch.pth')
+parser.add_argument('--load_a2l_C_name', type=str, default='examples/ckpt/ckpt_content_branch.pth') #ckpt_audio2landmark_c.pth')
 parser.add_argument('--load_G_name', type=str, default='examples/ckpt/ckpt_116_i2i_comb.pth') #ckpt_image2image.pth') #ckpt_i2i_finetune_150.pth') #c
 
 parser.add_argument('--amp_lip_x', type=float, default=2.)
 parser.add_argument('--amp_lip_y', type=float, default=2.)
-parser.add_argument('--amp_pos', type=float, default=1.)
+parser.add_argument('--amp_pos', type=float, default=.5)
 parser.add_argument('--reuse_train_emb_list', type=str, nargs='+', default=[]) #  ['iWeklsXc0H8']) #['45hn7-LXDX8']) #['E_kmpT-EfOg']) #'iWeklsXc0H8', '29k8RtSUjE0', '45hn7-LXDX8',
-# --reuse_train_emb_list 45hn7-LXDX8
-
 parser.add_argument('--add_audio_in', default=False, action='store_true')
 parser.add_argument('--comb_fan_awing', default=False, action='store_true')
 parser.add_argument('--output_folder', type=str, default='examples')
@@ -84,7 +81,7 @@
 
 
 ''' Additional manual adjustment to input face landmarks (slimmer lips and wider eyes) '''
-shape_3d[48:, 0] = (shape_3d[48:, 0] - np.mean(shape_3d[48:, 0])) * 0.95 + np.mean(shape_3d[48:, 0])
+# shape_3d[48:, 0] = (shape_3d[48:, 0] - np.mean(shape_3d[48:, 0])) * 0.95 + np.mean(shape_3d[48:, 0])
 shape_3d[49:54, 1] += 1.
 shape_3d[55:60, 1] -= 1.
 shape_3d[[37,38,43,44], 1] -=2

src/approaches/train_audio2landmark.py

@@ -97,7 +97,7 @@ def __init__(self, opt_parser, jpg_shape=None):
             print(key)
         print('====================================')
 
-    def __train_face_and_pos__(self, fls, aus, embs, face_id, smooth_win=31, close_mouth_ratio=.66):
+    def __train_face_and_pos__(self, fls, aus, embs, face_id, smooth_win=31, close_mouth_ratio=.99):
 
         fls_without_traj = fls[:, 0, :].detach().clone().requires_grad_(False)
 
@@ -107,7 +107,7 @@ def __train_face_and_pos__(self, fls, aus, embs, face_id, smooth_win=31, close_m
         baseline_face_id = face_id.detach()
 
         z = torch.tensor(torch.zeros(aus.shape[0], 128), requires_grad=False, dtype=torch.float).to(device)
-        fl_dis_pred, _, spk_encode = self.G(aus, embs * 3.0, face_id, fls_without_traj, z, add_z_spk=True)
+        fl_dis_pred, _, spk_encode = self.G(aus, embs * 3.0, face_id, fls_without_traj, z, add_z_spk=False)
 
         # ADD CONTENT
         from scipy.signal import savgol_filter
@@ -133,16 +133,30 @@ def __train_face_and_pos__(self, fls, aus, embs, face_id, smooth_win=31, close_m
 
         # ''' CALIBRATION '''
         baseline_pred_fls, _ = self.C(aus[:, 0:18, :], residual_face_id)
+        baseline_pred_fls = self.__calib_baseline_pred_fls__(baseline_pred_fls)
+        fl_dis_pred += baseline_pred_fls
+
+        return fl_dis_pred, face_id[0:1, :]
+
+    def __calib_baseline_pred_fls_old_(self, baseline_pred_fls, residual_face_id, aus):
         mean_face_id = torch.mean(baseline_pred_fls.detach(), dim=0, keepdim=True)
         residual_face_id -= mean_face_id.view(1, 204) * 1.
-        # ''' ======================== '''
-
-        baseline_pred_fls, _ = self.C(aus[:, 0:18, :], residual_face_id)
+        baseline_pred_fls, _ = self.C(aus, residual_face_id)
         baseline_pred_fls[:, 48 * 3::3] *= self.opt_parser.amp_lip_x  # mouth x
         baseline_pred_fls[:, 48 * 3 + 1::3] *= self.opt_parser.amp_lip_y  # mouth y
-        fl_dis_pred += baseline_pred_fls
-
-        return fl_dis_pred, face_id[0:1, :]
+        return baseline_pred_fls
+
+    def __calib_baseline_pred_fls__(self, baseline_pred_fls, ratio=0.5):
+        np_fl_dis_pred = baseline_pred_fls.detach().cpu().numpy()
+        K = int(np_fl_dis_pred.shape[0] * ratio)
+        for calib_i in range(204):
+            min_k_idx = np.argpartition(np_fl_dis_pred[:, calib_i], K)
+            m = np.mean(np_fl_dis_pred[min_k_idx[:K], calib_i])
+            np_fl_dis_pred[:, calib_i] = np_fl_dis_pred[:, calib_i] - m
+        baseline_pred_fls = torch.tensor(np_fl_dis_pred, requires_grad=False).to(device)
+        baseline_pred_fls[:, 48 * 3::3] *= self.opt_parser.amp_lip_x  # mouth x
+        baseline_pred_fls[:, 48 * 3 + 1::3] *= self.opt_parser.amp_lip_y  # mouth y
+        return baseline_pred_fls
 
     def __train_pass__(self, au_emb=None, centerize_face=False, no_y_rotation=False, vis_fls=False):
 

src/dataset/audio2landmark/audio2landmark_dataset.py

@@ -79,7 +79,7 @@ def my_collate_in_segments(self, batch):
         return fls, aus, embs
 
     def my_collate_in_segments_noemb(self, batch):
-        fls, aus, embs = [], [], []
+        fls, aus = [], []
         for fl, au in batch:
             fl_data, au_data = fl[0], au[0]
             assert (fl_data.shape[0] == au_data.shape[0])
@@ -229,7 +229,7 @@ def __init__(self, dump_dir, dump_name, num_window_frames, num_window_step, stat
         #     print('SAVE!')
 
 
-        au_mean_std = np.loadtxt('dataset/utils/MEAN_STD_AUTOVC_RETRAIN_MEL_AU.txt') # np.mean(self.au_data[0][0]), np.std(self.au_data[0][0])
+        au_mean_std = np.loadtxt('src/dataset/utils/MEAN_STD_AUTOVC_RETRAIN_MEL_AU.txt') # np.mean(self.au_data[0][0]), np.std(self.au_data[0][0])
         au_mean, au_std = au_mean_std[0:au_mean_std.shape[0]//2], au_mean_std[au_mean_std.shape[0]//2:]
 
         self.au_data = [((au - au_mean) / au_std, info) for au, info in self.au_data]

src/models/model_audio2landmark_speaker_aware.py

@@ -226,7 +226,7 @@ def forward(self, x, e_outputs, src_mask=None, trg_mask=None):
         return self.norm(x)
 
 
-class Audio2landmark_speaker_aware(nn.Module):
+class Audio2landmark_speaker_aware_old(nn.Module):
 
     def __init__(self, spk_emb_enc_size=128,
                  transformer_d_model=32, N=2, heads=2,
@@ -291,7 +291,7 @@ def __init__(self, spk_emb_enc_size=128,
         )
 
 
-    def forward(self, au, face_id, add_z_spk=False):
+    def forward(self, au, face_id):
 
         ''' original version '''
         # audio
@@ -321,6 +321,87 @@ def forward(self, au, face_id, add_z_spk=False):
         return fl_pred, pos_pred, face_id[0:1, :], None
 
 
+class Audio2landmark_speaker_aware(nn.Module):
+
+    def __init__(self, audio_feat_size=80, c_enc_hidden_size=256, num_layers=3, drop_out=0,
+                 spk_feat_size=256, spk_emb_enc_size=128, lstm_g_win_size=64, add_info_size=6,
+                 transformer_d_model=32, N=2, heads=2, z_size=128, audio_dim=256):
+        super(Audio2landmark_speaker_aware, self).__init__()
+
+        self.lstm_g_win_size = lstm_g_win_size
+        self.add_info_size = add_info_size
+        comb_mlp_size = c_enc_hidden_size * 2
+
+        self.audio_content_encoder = nn.LSTM(input_size=audio_feat_size,
+                                             hidden_size=c_enc_hidden_size,
+                                             num_layers=num_layers,
+                                             dropout=drop_out,
+                                             bidirectional=False,
+                                             batch_first=True)
+
+        self.use_audio_projection = not (audio_dim == c_enc_hidden_size)
+        if(self.use_audio_projection):
+            self.audio_projection = nn.Sequential(
+                nn.Linear(in_features=c_enc_hidden_size, out_features=256),
+                nn.LeakyReLU(0.02),
+                nn.Linear(256, 128),
+                nn.LeakyReLU(0.02),
+                nn.Linear(128, audio_dim),
+            )
+
+
+        ''' original version '''
+        self.spk_emb_encoder = nn.Sequential(
+            nn.Linear(in_features=spk_feat_size, out_features=256),
+            nn.LeakyReLU(0.02),
+            nn.Linear(256, 128),
+            nn.LeakyReLU(0.02),
+            nn.Linear(128, spk_emb_enc_size),
+        )
+
+        d_model = transformer_d_model * heads
+        N = N
+        heads = heads
+
+        self.encoder = Encoder(d_model, N, heads, in_size=audio_dim + spk_emb_enc_size + z_size)
+        self.decoder = Decoder(d_model, N, heads, in_size=204)
+        self.out = nn.Sequential(
+            nn.Linear(in_features=d_model + z_size, out_features=512),
+            nn.LeakyReLU(0.02),
+            nn.Linear(512, 256),
+            nn.LeakyReLU(0.02),
+            nn.Linear(256, 204),
+        )
+
+
+    def forward(self, au, emb, face_id, add_z_spk=False, another_emb=None):
+
+        # audio
+        audio_encode, (_, _) = self.audio_content_encoder(au)
+        audio_encode = audio_encode[:, -1, :]
+
+        if(self.use_audio_projection):
+            audio_encode = self.audio_projection(audio_encode)
+
+        # spk
+        spk_encode = self.spk_emb_encoder(emb)
+        if(add_z_spk):
+            z_spk = torch.tensor(torch.randn(spk_encode.shape)*0.01, requires_grad=False, dtype=torch.float).to(device)
+            spk_encode = spk_encode + z_spk
+
+        # comb
+        z = torch.tensor(torch.zeros(au.shape[0], 128), requires_grad=False, dtype=torch.float).to(device)
+        comb_encode = torch.cat((audio_encode, spk_encode, z), dim=1)
+        src_feat = comb_encode.unsqueeze(0)
+
+        e_outputs = self.encoder(src_feat)[0]
+
+        e_outputs = torch.cat((e_outputs, z), dim=1)
+
+        fl_pred = self.out(e_outputs)
+
+        return fl_pred, face_id[0:1, :], spk_encode
+
 
 
 def nopeak_mask(size):
@@ -344,23 +425,6 @@ def create_masks(src, trg):
     return src_mask, trg_mask
 
 
-class TalkingToon_spk2res_lstmgan_DL(nn.Module):
-    def __init__(self, comb_emb_size=256, input_size=6):
-        super(TalkingToon_spk2res_lstmgan_DL, self).__init__()
-
-        self.fl_D = nn.Sequential(
-            nn.Linear(in_features=FACE_ID_FEAT_SIZE, out_features=512),
-            nn.LeakyReLU(0.02),
-            nn.Linear(512, 256),
-            nn.LeakyReLU(0.02),
-            nn.Linear(256, 1),
-        )
-
-    def forward(self, feat):
-        d = self.fl_D(feat)
-        # d = torch.sigmoid(d)
-        return d
-
 
 class Transformer_DT(nn.Module):
     def __init__(self, transformer_d_model=32, N=2, heads=2, spk_emb_enc_size=128):
@@ -375,11 +439,11 @@ def __init__(self, transformer_d_model=32, N=2, heads=2, spk_emb_enc_size=128):
             nn.Linear(256, 1),
         )
 
-    def forward(self, fls, spk_emb, win_size=64, win_step=1):
+    def forward(self, fls, spk_emb, win_size=64, win_step=16):
         feat = torch.cat((fls, spk_emb), dim=1)
 
         win_size = feat.shape[0]-1 if feat.shape[0] <= win_size else win_size
-        D_input = [feat[i:i+win_size:win_step] for i in range(0, feat.shape[0]-win_size)]
+        D_input = [feat[i:i+win_size:win_step] for i in range(0, feat.shape[0]-win_size, win_step)]
         D_input = torch.stack(D_input, dim=0)
         D_output = self.encoder(D_input)
         D_output = torch.max(D_output, dim=1, keepdim=False)[0]

util/vis.py

@@ -102,7 +102,8 @@ def draw_curve(idx_list, color=(0, 255, 0), loop=False, lineWidth=linewidth):
 
 class Vis_old():
 
-    def __init__(self, run_name, pred_fl_filename, audio_filename, av_name='NAME', fps=100, frames=625, postfix='', root_dir=r'E:\Dataset\TalkingToon\Obama', ifsmooth=True):
+    def __init__(self, run_name, pred_fl_filename, audio_filename, av_name='NAME', fps=100, frames=625,
+                 postfix='', root_dir=r'E:\Dataset\TalkingToon\Obama', ifsmooth=True, rand_start=0):
 
         print(root_dir)
         self.src_dir = os.path.join(root_dir, r'nn_result/{}'.format(run_name))
@@ -140,13 +141,19 @@ def __init__(self, run_name, pred_fl_filename, audio_filename, av_name='NAME', f
         # out = out.overwrite_output().global_args('-loglevel', 'quiet')
         # out.run()
 
+        os.system('ffmpeg -y -loglevel error -i {} -ss {} {}'.format(
+            ain, rand_start/62.5,
+            os.path.join(self.src_dir, '{}_a_tmp.wav'.format(av_name))
+        ))
+
         os.system('ffmpeg -y -loglevel error -i {} -i {} -pix_fmt yuv420p -strict -2 -shortest {}'.format(
             os.path.join(self.src_dir, 'tmp.mp4'),
-            ain,
+            os.path.join(self.src_dir, '{}_a_tmp.wav'.format(av_name)),
             os.path.join(self.src_dir, '{}_av.mp4'.format(av_name))
         ))
 
         os.remove(os.path.join(self.src_dir, 'tmp.mp4'))
+        os.remove(os.path.join(self.src_dir, '{}_a_tmp.wav'.format(av_name)))
 
         # os.remove(os.path.join(self.src_dir, filename))
         # exit(0)