Spaces:
Sleeping
Sleeping
File size: 3,210 Bytes
579f772 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
#!/usr/bin/python
#-*- coding: utf-8 -*-
import torch
import numpy
import time, pdb, argparse, subprocess, pickle, os, glob
import cv2
from scipy import signal
# ==================== PARSE ARGUMENT ====================
parser = argparse.ArgumentParser(description = "SyncNet");
parser.add_argument('--data_dir', type=str, default='data/work', help='');
parser.add_argument('--videofile', type=str, default='', help='');
parser.add_argument('--reference', type=str, default='', help='');
parser.add_argument('--frame_rate', type=int, default=25, help='Frame rate');
opt = parser.parse_args();
setattr(opt,'avi_dir',os.path.join(opt.data_dir,'pyavi'))
setattr(opt,'tmp_dir',os.path.join(opt.data_dir,'pytmp'))
setattr(opt,'work_dir',os.path.join(opt.data_dir,'pywork'))
setattr(opt,'crop_dir',os.path.join(opt.data_dir,'pycrop'))
setattr(opt,'frames_dir',os.path.join(opt.data_dir,'pyframes'))
# ==================== LOAD FILES ====================
with open(os.path.join(opt.work_dir,opt.reference,'tracks.pckl'), 'rb') as fil:
tracks = pickle.load(fil, encoding='latin1')
with open(os.path.join(opt.work_dir,opt.reference,'activesd.pckl'), 'rb') as fil:
dists = pickle.load(fil, encoding='latin1')
flist = glob.glob(os.path.join(opt.frames_dir,opt.reference,'*.jpg'))
flist.sort()
# ==================== SMOOTH FACES ====================
faces = [[] for i in range(len(flist))]
for tidx, track in enumerate(tracks):
mean_dists = numpy.mean(numpy.stack(dists[tidx],1),1)
minidx = numpy.argmin(mean_dists,0)
minval = mean_dists[minidx]
fdist = numpy.stack([dist[minidx] for dist in dists[tidx]])
fdist = numpy.pad(fdist, (3,3), 'constant', constant_values=10)
fconf = numpy.median(mean_dists) - fdist
fconfm = signal.medfilt(fconf,kernel_size=9)
for fidx, frame in enumerate(track['track']['frame'].tolist()) :
faces[frame].append({'track': tidx, 'conf':fconfm[fidx], 's':track['proc_track']['s'][fidx], 'x':track['proc_track']['x'][fidx], 'y':track['proc_track']['y'][fidx]})
# ==================== ADD DETECTIONS TO VIDEO ====================
first_image = cv2.imread(flist[0])
fw = first_image.shape[1]
fh = first_image.shape[0]
fourcc = cv2.VideoWriter_fourcc(*'XVID')
vOut = cv2.VideoWriter(os.path.join(opt.avi_dir,opt.reference,'video_only.avi'), fourcc, opt.frame_rate, (fw,fh))
for fidx, fname in enumerate(flist):
image = cv2.imread(fname)
for face in faces[fidx]:
clr = max(min(face['conf']*25,255),0)
cv2.rectangle(image,(int(face['x']-face['s']),int(face['y']-face['s'])),(int(face['x']+face['s']),int(face['y']+face['s'])),(0,clr,255-clr),3)
cv2.putText(image,'Track %d, Conf %.3f'%(face['track'],face['conf']), (int(face['x']-face['s']),int(face['y']-face['s'])),cv2.FONT_HERSHEY_SIMPLEX,0.5,(255,255,255),2)
vOut.write(image)
print('Frame %d'%fidx)
vOut.release()
# ========== COMBINE AUDIO AND VIDEO FILES ==========
command = ("ffmpeg -y -i %s -i %s -c:v copy -c:a copy %s" % (os.path.join(opt.avi_dir,opt.reference,'video_only.avi'),os.path.join(opt.avi_dir,opt.reference,'audio.wav'),os.path.join(opt.avi_dir,opt.reference,'video_out.avi'))) #-async 1
output = subprocess.call(command, shell=True, stdout=None)
|