welher/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 rmtree13from latentsync.utils.util import check_model_and_download14 15 16# ==================== Get OFFSET ====================17 18# Video 25 FPS, Audio 16000HZ19 20 21def calc_pdist(feat1, feat2, vshift=10):22    win_size = vshift * 2 + 123 24    feat2p = torch.nn.functional.pad(feat2, (0, 0, vshift, vshift))25 26    dists = []27 28    for i in range(0, len(feat1)):29 30        dists.append(31            torch.nn.functional.pairwise_distance(feat1[[i], :].repeat(win_size, 1), feat2p[i : i + win_size, :])32        )33 34    return dists35 36 37# ==================== MAIN DEF ====================38 39 40class SyncNetEval(torch.nn.Module):41    def __init__(self, dropout=0, num_layers_in_fc_layers=1024, device="cpu"):42        super().__init__()43 44        self.__S__ = S(num_layers_in_fc_layers=num_layers_in_fc_layers).to(device)45        self.device = device46 47    def evaluate(self, video_path, temp_dir="temp", batch_size=20, vshift=15):48 49        self.__S__.eval()50 51        # ========== ==========52        # Convert files53        # ========== ==========54 55        if os.path.exists(temp_dir):56            rmtree(temp_dir)57 58        os.makedirs(temp_dir)59 60        # temp_video_path = os.path.join(temp_dir, "temp.mp4")61        # command = f"ffmpeg -loglevel error -nostdin -y -i {video_path} -vf scale='224:224' {temp_video_path}"62        # subprocess.call(command, shell=True)63 64        command = f"ffmpeg -loglevel error -nostdin -y -i {video_path} -f image2 {os.path.join(temp_dir, '%06d.jpg')}"65        subprocess.call(command, shell=True, stdout=None)66 67        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')}"68        subprocess.call(command, shell=True, stdout=None)69 70        # ========== ==========71        # Load video72        # ========== ==========73 74        images = []75 76        flist = glob.glob(os.path.join(temp_dir, "*.jpg"))77        flist.sort()78 79        for fname in flist:80            img_input = cv2.imread(fname)81            img_input = cv2.resize(img_input, (224, 224))  # HARD CODED, CHANGE BEFORE RELEASE82            images.append(img_input)83 84        im = numpy.stack(images, axis=3)85        im = numpy.expand_dims(im, axis=0)86        im = numpy.transpose(im, (0, 3, 4, 1, 2))87 88        imtv = torch.autograd.Variable(torch.from_numpy(im.astype(float)).float())89 90        # ========== ==========91        # Load audio92        # ========== ==========93 94        sample_rate, audio = wavfile.read(os.path.join(temp_dir, "audio.wav"))95        mfcc = zip(*python_speech_features.mfcc(audio, sample_rate))96        mfcc = numpy.stack([numpy.array(i) for i in mfcc])97 98        cc = numpy.expand_dims(numpy.expand_dims(mfcc, axis=0), axis=0)99        cct = torch.autograd.Variable(torch.from_numpy(cc.astype(float)).float())100 101        # ========== ==========102        # Check audio and video input length103        # ========== ==========104 105        # if (float(len(audio)) / 16000) != (float(len(images)) / 25):106        #     print(107        #         "WARNING: Audio (%.4fs) and video (%.4fs) lengths are different."108        #         % (float(len(audio)) / 16000, float(len(images)) / 25)109        #     )110 111        min_length = min(len(images), math.floor(len(audio) / 640))112 113        # ========== ==========114        # Generate video and audio feats115        # ========== ==========116 117        lastframe = min_length - 5118        im_feat = []119        cc_feat = []120 121        tS = time.time()122        for i in range(0, lastframe, batch_size):123 124            im_batch = [imtv[:, :, vframe : vframe + 5, :, :] for vframe in range(i, min(lastframe, i + batch_size))]125            im_in = torch.cat(im_batch, 0)126            im_out = self.__S__.forward_lip(im_in.to(self.device))127            im_feat.append(im_out.data.cpu())128 129            cc_batch = [130                cct[:, :, :, vframe * 4 : vframe * 4 + 20] for vframe in range(i, min(lastframe, i + batch_size))131            ]132            cc_in = torch.cat(cc_batch, 0)133            cc_out = self.__S__.forward_aud(cc_in.to(self.device))134            cc_feat.append(cc_out.data.cpu())135 136        im_feat = torch.cat(im_feat, 0)137        cc_feat = torch.cat(cc_feat, 0)138 139        # ========== ==========140        # Compute offset141        # ========== ==========142 143        dists = calc_pdist(im_feat, cc_feat, vshift=vshift)144        mean_dists = torch.mean(torch.stack(dists, 1), 1)145 146        min_dist, minidx = torch.min(mean_dists, 0)147 148        av_offset = vshift - minidx149        conf = torch.median(mean_dists) - min_dist150 151        fdist = numpy.stack([dist[minidx].numpy() for dist in dists])152        # fdist   = numpy.pad(fdist, (3,3), 'constant', constant_values=15)153        fconf = torch.median(mean_dists).numpy() - fdist154        framewise_conf = signal.medfilt(fconf, kernel_size=9)155 156        # numpy.set_printoptions(formatter={"float": "{: 0.3f}".format})157        rmtree(temp_dir)158        return av_offset.item(), min_dist.item(), conf.item()159 160    def extract_feature(self, opt, videofile):161 162        self.__S__.eval()163 164        # ========== ==========165        # Load video166        # ========== ==========167        cap = cv2.VideoCapture(videofile)168 169        frame_num = 1170        images = []171        while frame_num:172            frame_num += 1173            ret, image = cap.read()174            if ret == 0:175                break176 177            images.append(image)178 179        im = numpy.stack(images, axis=3)180        im = numpy.expand_dims(im, axis=0)181        im = numpy.transpose(im, (0, 3, 4, 1, 2))182 183        imtv = torch.autograd.Variable(torch.from_numpy(im.astype(float)).float())184 185        # ========== ==========186        # Generate video feats187        # ========== ==========188 189        lastframe = len(images) - 4190        im_feat = []191 192        tS = time.time()193        for i in range(0, lastframe, opt.batch_size):194 195            im_batch = [196                imtv[:, :, vframe : vframe + 5, :, :] for vframe in range(i, min(lastframe, i + opt.batch_size))197            ]198            im_in = torch.cat(im_batch, 0)199            im_out = self.__S__.forward_lipfeat(im_in.to(self.device))200            im_feat.append(im_out.data.cpu())201 202        im_feat = torch.cat(im_feat, 0)203 204        # ========== ==========205        # Compute offset206        # ========== ==========207 208        print("Compute time %.3f sec." % (time.time() - tS))209 210        return im_feat211 212    def loadParameters(self, path):213        check_model_and_download(path)214        loaded_state = torch.load(path, map_location=lambda storage, loc: storage, weights_only=True)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