RabbitRUI/ruispace

1from tqdm import tqdm2import torch3from torch import nn4 5 6class Audio2Exp(nn.Module):7    def __init__(self, netG, cfg, device, prepare_training_loss=False):8        super(Audio2Exp, self).__init__()9        self.cfg = cfg10        self.device = device11        self.netG = netG.to(device)12 13    def test(self, batch):14 15        mel_input = batch['indiv_mels']                         # bs T 1 80 1616        bs = mel_input.shape[0]17        T = mel_input.shape[1]18 19        exp_coeff_pred = []20 21        for i in tqdm(range(0, T, 10),'audio2exp:'): # every 10 frames22            23            current_mel_input = mel_input[:,i:i+10]24 25            #ref = batch['ref'][:, :, :64].repeat((1,current_mel_input.shape[1],1))           #bs T 6426            ref = batch['ref'][:, :, :64][:, i:i+10]27            ratio = batch['ratio_gt'][:, i:i+10]                               #bs T28 29            audiox = current_mel_input.view(-1, 1, 80, 16)                  # bs*T 1 80 1630 31            curr_exp_coeff_pred  = self.netG(audiox, ref, ratio)         # bs T 64 32 33            exp_coeff_pred += [curr_exp_coeff_pred]34 35        # BS x T x 6436        results_dict = {37            'exp_coeff_pred': torch.cat(exp_coeff_pred, axis=1)38            }39        return results_dict40 41 42