SyncNetInstance.py

#!/usr/bin/env python3
#-*- coding: utf-8 -*-
# Video 25 FPS, Audio 16000HZ

import torch
import numpy
import time, pdb, argparse, subprocess, os, math, glob, logging
import cv2
import python_speech_features

from scipy import signal
from scipy.io import wavfile
from SyncNetModel import *
from shutil import rmtree

logger = logging.getLogger(__name__)


# ==================== Get OFFSET ====================

def calc_pdist(feat1, feat2, vshift=10):

 win_size = vshift*2+1

 feat2p = torch.nn.functional.pad(feat2,(0,0,vshift,vshift))

 dists = []

 for i in range(0,len(feat1)):

 dists.append(torch.nn.functional.pairwise_distance(feat1[[i],:].repeat(win_size, 1), feat2p[i:i+win_size,:]))

 return dists

# ==================== MAIN DEF ====================

class SyncNetInstance(torch.nn.Module):

 def __init__(self, dropout = 0, num_layers_in_fc_layers = 1024, device=None):
 super().__init__()

 self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu')
 logger.info('Using device: %s', self.device)
 self.__S__ = S(num_layers_in_fc_layers = num_layers_in_fc_layers).to(self.device)

 def evaluate(self, opt, videofile):

 self.__S__.eval()

 # ========== ==========
 # Convert files
 # ========== ==========

 if os.path.exists(os.path.join(opt.tmp_dir,opt.reference)):
 rmtree(os.path.join(opt.tmp_dir,opt.reference))

 os.makedirs(os.path.join(opt.tmp_dir,opt.reference))

 logger.info('Extracting video frames from %s', videofile)
 command = ["ffmpeg", "-y", "-loglevel", "error", "-i", videofile, "-threads", "1", "-f", "image2",
 os.path.join(opt.tmp_dir, opt.reference, '%06d.jpg')]
 subprocess.run(command, check=True)

 logger.info('Extracting audio from %s', videofile)
 command = ["ffmpeg", "-y", "-loglevel", "error", "-i", videofile, "-async", "1", "-ac", "1", "-vn",
 "-acodec", "pcm_s16le", "-ar", "16000",
 os.path.join(opt.tmp_dir, opt.reference, 'audio.wav')]
 subprocess.run(command, check=True)

 # ========== ==========
 # Load video
 # ========== ==========

 images = []

 flist = glob.glob(os.path.join(opt.tmp_dir,opt.reference,'*.jpg'))
 flist.sort()

 for fname in flist:
 images.append(cv2.imread(fname))

 im = numpy.stack(images,axis=3)
 im = numpy.expand_dims(im,axis=0)
 im = numpy.transpose(im,(0,3,4,1,2))

 imtv = torch.from_numpy(im.astype(float)).float()

 # ========== ==========
 # Load audio
 # ========== ==========

 sample_rate, audio = wavfile.read(os.path.join(opt.tmp_dir,opt.reference,'audio.wav'))
 mfcc = zip(*python_speech_features.mfcc(audio,sample_rate))
 mfcc = numpy.stack([numpy.array(i) for i in mfcc])

 cc = numpy.expand_dims(numpy.expand_dims(mfcc,axis=0),axis=0)
 cct = torch.from_numpy(cc.astype(float)).float()

 # ========== ==========
 # Check audio and video input length
 # ========== ==========

 if (float(len(audio))/16000) != (float(len(images))/25) :
 logger.warning("Audio (%.4fs) and video (%.4fs) lengths are different.",float(len(audio))/16000,float(len(images))/25)

 min_length = min(len(images),math.floor(len(audio)/640))

 # ========== ==========
 # Generate video and audio feats
 # ========== ==========

 lastframe = min_length-5
 im_feat = []
 cc_feat = []

 tS = time.time()
 for i in range(0,lastframe,opt.batch_size):

 im_batch = [ imtv[:,:,vframe:vframe+5,:,:] for vframe in range(i,min(lastframe,i+opt.batch_size)) ]
 im_in = torch.cat(im_batch,0)
 im_out = self.__S__.forward_lip(im_in.to(self.device))
 im_feat.append(im_out.data.cpu())

 cc_batch = [ cct[:,:,:,vframe*4:vframe*4+20] for vframe in range(i,min(lastframe,i+opt.batch_size)) ]
 cc_in = torch.cat(cc_batch,0)
 cc_out = self.__S__.forward_aud(cc_in.to(self.device))
 cc_feat.append(cc_out.data.cpu())

 im_feat = torch.cat(im_feat,0)
 cc_feat = torch.cat(cc_feat,0)

 # ========== ==========
 # Compute offset
 # ========== ==========

 logger.info('Compute time %.3f sec.', time.time()-tS)

 dists = calc_pdist(im_feat,cc_feat,vshift=opt.vshift)
 mdist = torch.mean(torch.stack(dists,1),1)

 minval, minidx = torch.min(mdist,0)

 offset = opt.vshift-minidx
 conf = torch.median(mdist) - minval

 fdist = numpy.stack([dist[minidx].numpy() for dist in dists])
 # fdist = numpy.pad(fdist, (3,3), 'constant', constant_values=15)
 fconf = torch.median(mdist).numpy() - fdist
 fconfm = signal.medfilt(fconf,kernel_size=9)

 numpy.set_printoptions(formatter={'float': '{: 0.3f}'.format})
 logger.info('Framewise conf: ')
 logger.info(fconfm)
 logger.info('AV offset: \t%d', offset.item())
 logger.info('Min dist: \t%.3f', minval.item())
 logger.info('Confidence: \t%.3f', conf.item())

 dists_npy = numpy.array([ dist.numpy() for dist in dists ])
 return offset.numpy(), conf.numpy(), dists_npy

 def extract_feature(self, opt, videofile):

 self.__S__.eval()

 # ========== ==========
 # Load video
 # ========== ==========
 cap = cv2.VideoCapture(videofile)

 frame_num = 1
 images = []
 while frame_num:
 frame_num += 1
 ret, image = cap.read()
 if ret == 0:
 break

 images.append(image)

 im = numpy.stack(images,axis=3)
 im = numpy.expand_dims(im,axis=0)
 im = numpy.transpose(im,(0,3,4,1,2))

 imtv = torch.from_numpy(im.astype(float)).float()

 # ========== ==========
 # Generate video feats
 # ========== ==========

 lastframe = len(images)-4
 im_feat = []

 tS = time.time()
 for i in range(0,lastframe,opt.batch_size):

 im_batch = [ imtv[:,:,vframe:vframe+5,:,:] for vframe in range(i,min(lastframe,i+opt.batch_size)) ]
 im_in = torch.cat(im_batch,0)
 im_out = self.__S__.forward_lipfeat(im_in.to(self.device))
 im_feat.append(im_out.data.cpu())

 im_feat = torch.cat(im_feat,0)

 # ========== ==========
 # Compute offset
 # ========== ==========

 logger.info('Compute time %.3f sec.', time.time()-tS)

 return im_feat


 def loadParameters(self, path):
 loaded_state = torch.load(path, map_location=lambda storage, loc: storage, weights_only=True)

 self_state = self.__S__.state_dict()

 for name, param in loaded_state.items():

 self_state[name].copy_(param)