Francke/LatentSync

1# Adapted from https://github.com/joonson/syncnet_python/blob/master/SyncNetInstance.py2 3import torch4import numpy5import time, pdb, argparse, subprocess, os, math, glob6import cv27import python_speech_features8 9from scipy import signal10from scipy.io import wavfile11from .syncnet import S12from shutil import rmtree13 14 15# ==================== Get OFFSET ====================16 17# Video 25 FPS, Audio 16000HZ18 19 20def calc_pdist(feat1, feat2, vshift=10):21    win_size = vshift * 2 + 122 23    feat2p = torch.nn.functional.pad(feat2, (0, 0, vshift, vshift))24 25    dists = []26 27    for i in range(0, len(feat1)):28 29        dists.append(30            torch.nn.functional.pairwise_distance(feat1[[i], :].repeat(win_size, 1), feat2p[i : i + win_size, :])31        )32 33    return dists34 35 36# ==================== MAIN DEF ====================37 38 39class SyncNetEval(torch.nn.Module):40    def __init__(self, dropout=0, num_layers_in_fc_layers=1024, device="cpu"):41        super().__init__()42 43        self.__S__ = S(num_layers_in_fc_layers=num_layers_in_fc_layers).to(device)44        self.device = device45 46    def evaluate(self, video_path, temp_dir="temp", batch_size=20, vshift=15):47 48        self.__S__.eval()49 50        # ========== ==========51        # Convert files52        # ========== ==========53 54        if os.path.exists(temp_dir):55            rmtree(temp_dir)56 57        os.makedirs(temp_dir)58 59        # temp_video_path = os.path.join(temp_dir, "temp.mp4")60        # command = f"ffmpeg -loglevel error -nostdin -y -i {video_path} -vf scale='224:224' {temp_video_path}"61        # subprocess.call(command, shell=True)62 63        command = (64            f"ffmpeg -loglevel error -nostdin -y -i {video_path} -f image2 {os.path.join(temp_dir, '%06d.jpg')}"65        )66        subprocess.call(command, shell=True, stdout=None)67 68        command = f"ffmpeg -loglevel error -nostdin -y -i {video_path} -async 1 -ac 1 -vn -acodec pcm_s16le -ar 16000 {os.path.join(temp_dir, 'audio.wav')}"69        subprocess.call(command, shell=True, stdout=None)70 71        # ========== ==========72        # Load video73        # ========== ==========74 75        images = []76 77        flist = glob.glob(os.path.join(temp_dir, "*.jpg"))78        flist.sort()79 80        for fname in flist:81            img_input = cv2.imread(fname)82            img_input = cv2.resize(img_input, (224, 224))  # HARD CODED, CHANGE BEFORE RELEASE83            images.append(img_input)84 85        im = numpy.stack(images, axis=3)86        im = numpy.expand_dims(im, axis=0)87        im = numpy.transpose(im, (0, 3, 4, 1, 2))88 89        imtv = torch.autograd.Variable(torch.from_numpy(im.astype(float)).float())90 91        # ========== ==========92        # Load audio93        # ========== ==========94 95        sample_rate, audio = wavfile.read(os.path.join(temp_dir, "audio.wav"))96        mfcc = zip(*python_speech_features.mfcc(audio, sample_rate))97        mfcc = numpy.stack([numpy.array(i) for i in mfcc])98 99        cc = numpy.expand_dims(numpy.expand_dims(mfcc, axis=0), axis=0)100        cct = torch.autograd.Variable(torch.from_numpy(cc.astype(float)).float())101 102        # ========== ==========103        # Check audio and video input length104        # ========== ==========105 106        # if (float(len(audio)) / 16000) != (float(len(images)) / 25):107        #     print(108        #         "WARNING: Audio (%.4fs) and video (%.4fs) lengths are different."109        #         % (float(len(audio)) / 16000, float(len(images)) / 25)110        #     )111 112        min_length = min(len(images), math.floor(len(audio) / 640))113 114        # ========== ==========115        # Generate video and audio feats116        # ========== ==========117 118        lastframe = min_length - 5119        im_feat = []120        cc_feat = []121 122        tS = time.time()123        for i in range(0, lastframe, batch_size):124 125            im_batch = [imtv[:, :, vframe : vframe + 5, :, :] for vframe in range(i, min(lastframe, i + batch_size))]126            im_in = torch.cat(im_batch, 0)127            im_out = self.__S__.forward_lip(im_in.to(self.device))128            im_feat.append(im_out.data.cpu())129 130            cc_batch = [131                cct[:, :, :, vframe * 4 : vframe * 4 + 20] for vframe in range(i, min(lastframe, i + batch_size))132            ]133            cc_in = torch.cat(cc_batch, 0)134            cc_out = self.__S__.forward_aud(cc_in.to(self.device))135            cc_feat.append(cc_out.data.cpu())136 137        im_feat = torch.cat(im_feat, 0)138        cc_feat = torch.cat(cc_feat, 0)139 140        # ========== ==========141        # Compute offset142        # ========== ==========143 144        dists = calc_pdist(im_feat, cc_feat, vshift=vshift)145        mean_dists = torch.mean(torch.stack(dists, 1), 1)146 147        min_dist, minidx = torch.min(mean_dists, 0)148 149        av_offset = vshift - minidx150        conf = torch.median(mean_dists) - min_dist151 152        fdist = numpy.stack([dist[minidx].numpy() for dist in dists])153        # fdist   = numpy.pad(fdist, (3,3), 'constant', constant_values=15)154        fconf = torch.median(mean_dists).numpy() - fdist155        framewise_conf = signal.medfilt(fconf, kernel_size=9)156 157        # numpy.set_printoptions(formatter={"float": "{: 0.3f}".format})158        rmtree(temp_dir)159        return av_offset.item(), min_dist.item(), conf.item()160 161    def extract_feature(self, opt, videofile):162 163        self.__S__.eval()164 165        # ========== ==========166        # Load video167        # ========== ==========168        cap = cv2.VideoCapture(videofile)169 170        frame_num = 1171        images = []172        while frame_num:173            frame_num += 1174            ret, image = cap.read()175            if ret == 0:176                break177 178            images.append(image)179 180        im = numpy.stack(images, axis=3)181        im = numpy.expand_dims(im, axis=0)182        im = numpy.transpose(im, (0, 3, 4, 1, 2))183 184        imtv = torch.autograd.Variable(torch.from_numpy(im.astype(float)).float())185 186        # ========== ==========187        # Generate video feats188        # ========== ==========189 190        lastframe = len(images) - 4191        im_feat = []192 193        tS = time.time()194        for i in range(0, lastframe, opt.batch_size):195 196            im_batch = [197                imtv[:, :, vframe : vframe + 5, :, :] for vframe in range(i, min(lastframe, i + opt.batch_size))198            ]199            im_in = torch.cat(im_batch, 0)200            im_out = self.__S__.forward_lipfeat(im_in.to(self.device))201            im_feat.append(im_out.data.cpu())202 203        im_feat = torch.cat(im_feat, 0)204 205        # ========== ==========206        # Compute offset207        # ========== ==========208 209        print("Compute time %.3f sec." % (time.time() - tS))210 211        return im_feat212 213    def loadParameters(self, path):214        loaded_state = torch.load(path, map_location=lambda storage, loc: storage)215 216        self_state = self.__S__.state_dict()217 218        for name, param in loaded_state.items():219 220            self_state[name].copy_(param)221