362 lines
16 KiB
Python
362 lines
16 KiB
Python
import torch
|
|
|
|
import sys, time, os, tqdm, torch, argparse, glob, subprocess, warnings, cv2, pickle, pdb, math, python_speech_features
|
|
import numpy as np
|
|
from scipy import signal
|
|
from shutil import rmtree
|
|
from scipy.io import wavfile
|
|
from scipy.interpolate import interp1d
|
|
from sklearn.metrics import accuracy_score, f1_score
|
|
import soundfile as sf
|
|
|
|
from scenedetect.video_manager import VideoManager
|
|
from scenedetect.scene_manager import SceneManager
|
|
from scenedetect.frame_timecode import FrameTimecode
|
|
from scenedetect.stats_manager import StatsManager
|
|
from scenedetect.detectors import ContentDetector
|
|
|
|
from models.av_mossformer2_tse.faceDetector.s3fd import S3FD
|
|
|
|
from .decode import decode_one_audio_AV_MossFormer2_TSE_16K
|
|
|
|
|
|
|
|
def process_tse(args, model, device, data_reader, output_wave_dir):
|
|
video_args = args_param()
|
|
video_args.model = model
|
|
video_args.device = device
|
|
video_args.sampling_rate = args.sampling_rate
|
|
args.device = device
|
|
assert args.sampling_rate == 16000
|
|
with torch.no_grad():
|
|
for videoPath in data_reader: # Loop over all video samples
|
|
savFolder = videoPath.split(os.path.sep)[-1]
|
|
video_args.savePath = f'{output_wave_dir}/{savFolder.split(".")[0]}/'
|
|
video_args.videoPath = videoPath
|
|
main(video_args, args)
|
|
|
|
|
|
|
|
def args_param():
|
|
warnings.filterwarnings("ignore")
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--nDataLoaderThread', type=int, default=10, help='Number of workers')
|
|
parser.add_argument('--facedetScale', type=float, default=0.25, help='Scale factor for face detection, the frames will be scale to 0.25 orig')
|
|
parser.add_argument('--minTrack', type=int, default=50, help='Number of min frames for each shot')
|
|
parser.add_argument('--numFailedDet', type=int, default=10, help='Number of missed detections allowed before tracking is stopped')
|
|
parser.add_argument('--minFaceSize', type=int, default=1, help='Minimum face size in pixels')
|
|
parser.add_argument('--cropScale', type=float, default=0.40, help='Scale bounding box')
|
|
parser.add_argument('--start', type=int, default=0, help='The start time of the video')
|
|
parser.add_argument('--duration', type=int, default=0, help='The duration of the video, when set as 0, will extract the whole video')
|
|
video_args = parser.parse_args()
|
|
return video_args
|
|
|
|
|
|
# Main function
|
|
def main(video_args, args):
|
|
# Initialization
|
|
video_args.pyaviPath = os.path.join(video_args.savePath, 'py_video')
|
|
video_args.pyframesPath = os.path.join(video_args.savePath, 'pyframes')
|
|
video_args.pyworkPath = os.path.join(video_args.savePath, 'pywork')
|
|
video_args.pycropPath = os.path.join(video_args.savePath, 'py_faceTracks')
|
|
if os.path.exists(video_args.savePath):
|
|
rmtree(video_args.savePath)
|
|
os.makedirs(video_args.pyaviPath, exist_ok = True) # The path for the input video, input audio, output video
|
|
os.makedirs(video_args.pyframesPath, exist_ok = True) # Save all the video frames
|
|
os.makedirs(video_args.pyworkPath, exist_ok = True) # Save the results in this process by the pckl method
|
|
os.makedirs(video_args.pycropPath, exist_ok = True) # Save the detected face clips (audio+video) in this process
|
|
|
|
# Extract video
|
|
video_args.videoFilePath = os.path.join(video_args.pyaviPath, 'video.avi')
|
|
# If duration did not set, extract the whole video, otherwise extract the video from 'video_args.start' to 'video_args.start + video_args.duration'
|
|
if video_args.duration == 0:
|
|
command = ("ffmpeg -y -i %s -qscale:v 2 -threads %d -async 1 -r 25 %s -loglevel panic" % \
|
|
(video_args.videoPath, video_args.nDataLoaderThread, video_args.videoFilePath))
|
|
else:
|
|
command = ("ffmpeg -y -i %s -qscale:v 2 -threads %d -ss %.3f -to %.3f -async 1 -r 25 %s -loglevel panic" % \
|
|
(video_args.videoPath, video_args.nDataLoaderThread, video_args.start, video_args.start + video_args.duration, video_args.videoFilePath))
|
|
subprocess.call(command, shell=True, stdout=None)
|
|
sys.stderr.write(time.strftime("%Y-%m-%d %H:%M:%S") + " Extract the video and save in %s \r\n" %(video_args.videoFilePath))
|
|
|
|
# Extract audio
|
|
video_args.audioFilePath = os.path.join(video_args.pyaviPath, 'audio.wav')
|
|
command = ("ffmpeg -y -i %s -qscale:a 0 -ac 1 -vn -threads %d -ar 16000 %s -loglevel panic" % \
|
|
(video_args.videoFilePath, video_args.nDataLoaderThread, video_args.audioFilePath))
|
|
subprocess.call(command, shell=True, stdout=None)
|
|
sys.stderr.write(time.strftime("%Y-%m-%d %H:%M:%S") + " Extract the audio and save in %s \r\n" %(video_args.audioFilePath))
|
|
|
|
# Extract the video frames
|
|
command = ("ffmpeg -y -i %s -qscale:v 2 -threads %d -f image2 %s -loglevel panic" % \
|
|
(video_args.videoFilePath, video_args.nDataLoaderThread, os.path.join(video_args.pyframesPath, '%06d.jpg')))
|
|
subprocess.call(command, shell=True, stdout=None)
|
|
sys.stderr.write(time.strftime("%Y-%m-%d %H:%M:%S") + " Extract the frames and save in %s \r\n" %(video_args.pyframesPath))
|
|
|
|
# Scene detection for the video frames
|
|
scene = scene_detect(video_args)
|
|
sys.stderr.write(time.strftime("%Y-%m-%d %H:%M:%S") + " Scene detection and save in %s \r\n" %(video_args.pyworkPath))
|
|
|
|
# Face detection for the video frames
|
|
faces = inference_video(video_args)
|
|
sys.stderr.write(time.strftime("%Y-%m-%d %H:%M:%S") + " Face detection and save in %s \r\n" %(video_args.pyworkPath))
|
|
|
|
# Face tracking
|
|
allTracks, vidTracks = [], []
|
|
for shot in scene:
|
|
if shot[1].frame_num - shot[0].frame_num >= video_args.minTrack: # Discard the shot frames less than minTrack frames
|
|
allTracks.extend(track_shot(video_args, faces[shot[0].frame_num:shot[1].frame_num])) # 'frames' to present this tracks' timestep, 'bbox' presents the location of the faces
|
|
sys.stderr.write(time.strftime("%Y-%m-%d %H:%M:%S") + " Face track and detected %d tracks \r\n" %len(allTracks))
|
|
|
|
# Face clips cropping
|
|
for ii, track in tqdm.tqdm(enumerate(allTracks), total = len(allTracks)):
|
|
vidTracks.append(crop_video(video_args, track, os.path.join(video_args.pycropPath, '%05d'%ii)))
|
|
savePath = os.path.join(video_args.pyworkPath, 'tracks.pckl')
|
|
with open(savePath, 'wb') as fil:
|
|
pickle.dump(vidTracks, fil)
|
|
sys.stderr.write(time.strftime("%Y-%m-%d %H:%M:%S") + " Face Crop and saved in %s tracks \r\n" %video_args.pycropPath)
|
|
fil = open(savePath, 'rb')
|
|
vidTracks = pickle.load(fil)
|
|
fil.close()
|
|
|
|
# AVSE
|
|
files = glob.glob("%s/*.avi"%video_args.pycropPath)
|
|
files.sort()
|
|
|
|
est_sources = evaluate_network(files, video_args, args)
|
|
|
|
visualization(vidTracks, est_sources, video_args)
|
|
|
|
# combine files in pycrop
|
|
for idx, file in enumerate(files):
|
|
print(file)
|
|
command = f"ffmpeg -i {file} {file[:-9]}orig_{idx}.mp4 ;"
|
|
command += f"rm {file} ;"
|
|
command += f"rm {file.replace('.avi', '.wav')} ;"
|
|
|
|
command += f"ffmpeg -i {file[:-9]}orig_{idx}.mp4 -i {file[:-9]}est_{idx}.wav -c:v copy -map 0:v:0 -map 1:a:0 -shortest {file[:-9]}est_{idx}.mp4 ;"
|
|
# command += f"rm {file[:-9]}est_{idx}.wav ;"
|
|
|
|
output = subprocess.call(command, shell=True, stdout=None)
|
|
|
|
rmtree(video_args.pyworkPath)
|
|
rmtree(video_args.pyframesPath)
|
|
|
|
|
|
|
|
|
|
def scene_detect(video_args):
|
|
# CPU: Scene detection, output is the list of each shot's time duration
|
|
videoManager = VideoManager([video_args.videoFilePath])
|
|
statsManager = StatsManager()
|
|
sceneManager = SceneManager(statsManager)
|
|
sceneManager.add_detector(ContentDetector())
|
|
baseTimecode = videoManager.get_base_timecode()
|
|
videoManager.set_downscale_factor()
|
|
videoManager.start()
|
|
sceneManager.detect_scenes(frame_source = videoManager)
|
|
sceneList = sceneManager.get_scene_list(baseTimecode)
|
|
savePath = os.path.join(video_args.pyworkPath, 'scene.pckl')
|
|
if sceneList == []:
|
|
sceneList = [(videoManager.get_base_timecode(),videoManager.get_current_timecode())]
|
|
with open(savePath, 'wb') as fil:
|
|
pickle.dump(sceneList, fil)
|
|
sys.stderr.write('%s - scenes detected %d\n'%(video_args.videoFilePath, len(sceneList)))
|
|
return sceneList
|
|
|
|
def inference_video(video_args):
|
|
# GPU: Face detection, output is the list contains the face location and score in this frame
|
|
DET = S3FD(device=video_args.device)
|
|
flist = glob.glob(os.path.join(video_args.pyframesPath, '*.jpg'))
|
|
flist.sort()
|
|
dets = []
|
|
for fidx, fname in enumerate(flist):
|
|
image = cv2.imread(fname)
|
|
imageNumpy = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
bboxes = DET.detect_faces(imageNumpy, conf_th=0.9, scales=[video_args.facedetScale])
|
|
dets.append([])
|
|
for bbox in bboxes:
|
|
dets[-1].append({'frame':fidx, 'bbox':(bbox[:-1]).tolist(), 'conf':bbox[-1]}) # dets has the frames info, bbox info, conf info
|
|
sys.stderr.write('%s-%05d; %d dets\r' % (video_args.videoFilePath, fidx, len(dets[-1])))
|
|
savePath = os.path.join(video_args.pyworkPath,'faces.pckl')
|
|
with open(savePath, 'wb') as fil:
|
|
pickle.dump(dets, fil)
|
|
return dets
|
|
|
|
def bb_intersection_over_union(boxA, boxB, evalCol = False):
|
|
# CPU: IOU Function to calculate overlap between two image
|
|
xA = max(boxA[0], boxB[0])
|
|
yA = max(boxA[1], boxB[1])
|
|
xB = min(boxA[2], boxB[2])
|
|
yB = min(boxA[3], boxB[3])
|
|
interArea = max(0, xB - xA) * max(0, yB - yA)
|
|
boxAArea = (boxA[2] - boxA[0]) * (boxA[3] - boxA[1])
|
|
boxBArea = (boxB[2] - boxB[0]) * (boxB[3] - boxB[1])
|
|
if evalCol == True:
|
|
iou = interArea / float(boxAArea)
|
|
else:
|
|
iou = interArea / float(boxAArea + boxBArea - interArea)
|
|
return iou
|
|
|
|
def track_shot(video_args, sceneFaces):
|
|
# CPU: Face tracking
|
|
iouThres = 0.5 # Minimum IOU between consecutive face detections
|
|
tracks = []
|
|
while True:
|
|
track = []
|
|
for frameFaces in sceneFaces:
|
|
for face in frameFaces:
|
|
if track == []:
|
|
track.append(face)
|
|
frameFaces.remove(face)
|
|
elif face['frame'] - track[-1]['frame'] <= video_args.numFailedDet:
|
|
iou = bb_intersection_over_union(face['bbox'], track[-1]['bbox'])
|
|
if iou > iouThres:
|
|
track.append(face)
|
|
frameFaces.remove(face)
|
|
continue
|
|
else:
|
|
break
|
|
if track == []:
|
|
break
|
|
elif len(track) > video_args.minTrack:
|
|
frameNum = np.array([ f['frame'] for f in track ])
|
|
bboxes = np.array([np.array(f['bbox']) for f in track])
|
|
frameI = np.arange(frameNum[0],frameNum[-1]+1)
|
|
bboxesI = []
|
|
for ij in range(0,4):
|
|
interpfn = interp1d(frameNum, bboxes[:,ij])
|
|
bboxesI.append(interpfn(frameI))
|
|
bboxesI = np.stack(bboxesI, axis=1)
|
|
if max(np.mean(bboxesI[:,2]-bboxesI[:,0]), np.mean(bboxesI[:,3]-bboxesI[:,1])) > video_args.minFaceSize:
|
|
tracks.append({'frame':frameI,'bbox':bboxesI})
|
|
return tracks
|
|
|
|
def crop_video(video_args, track, cropFile):
|
|
# CPU: crop the face clips
|
|
flist = glob.glob(os.path.join(video_args.pyframesPath, '*.jpg')) # Read the frames
|
|
flist.sort()
|
|
vOut = cv2.VideoWriter(cropFile + 't.avi', cv2.VideoWriter_fourcc(*'XVID'), 25, (224,224))# Write video
|
|
dets = {'x':[], 'y':[], 's':[]}
|
|
for det in track['bbox']: # Read the tracks
|
|
dets['s'].append(max((det[3]-det[1]), (det[2]-det[0]))/2)
|
|
dets['y'].append((det[1]+det[3])/2) # crop center x
|
|
dets['x'].append((det[0]+det[2])/2) # crop center y
|
|
dets['s'] = signal.medfilt(dets['s'], kernel_size=13) # Smooth detections
|
|
dets['x'] = signal.medfilt(dets['x'], kernel_size=13)
|
|
dets['y'] = signal.medfilt(dets['y'], kernel_size=13)
|
|
for fidx, frame in enumerate(track['frame']):
|
|
cs = video_args.cropScale
|
|
bs = dets['s'][fidx] # Detection box size
|
|
bsi = int(bs * (1 + 2 * cs)) # Pad videos by this amount
|
|
image = cv2.imread(flist[frame])
|
|
frame = np.pad(image, ((bsi,bsi), (bsi,bsi), (0, 0)), 'constant', constant_values=(110, 110))
|
|
my = dets['y'][fidx] + bsi # BBox center Y
|
|
mx = dets['x'][fidx] + bsi # BBox center X
|
|
face = frame[int(my-bs):int(my+bs*(1+2*cs)),int(mx-bs*(1+cs)):int(mx+bs*(1+cs))]
|
|
vOut.write(cv2.resize(face, (224, 224)))
|
|
audioTmp = cropFile + '.wav'
|
|
audioStart = (track['frame'][0]) / 25
|
|
audioEnd = (track['frame'][-1]+1) / 25
|
|
vOut.release()
|
|
command = ("ffmpeg -y -i %s -async 1 -ac 1 -vn -acodec pcm_s16le -ar 16000 -threads %d -ss %.3f -to %.3f %s -loglevel panic" % \
|
|
(video_args.audioFilePath, video_args.nDataLoaderThread, audioStart, audioEnd, audioTmp))
|
|
output = subprocess.call(command, shell=True, stdout=None) # Crop audio file
|
|
_, audio = wavfile.read(audioTmp)
|
|
command = ("ffmpeg -y -i %st.avi -i %s -threads %d -c:v copy -c:a copy %s.avi -loglevel panic" % \
|
|
(cropFile, audioTmp, video_args.nDataLoaderThread, cropFile)) # Combine audio and video file
|
|
output = subprocess.call(command, shell=True, stdout=None)
|
|
os.remove(cropFile + 't.avi')
|
|
return {'track':track, 'proc_track':dets}
|
|
|
|
|
|
def evaluate_network(files, video_args, args):
|
|
|
|
est_sources = []
|
|
for file in tqdm.tqdm(files, total = len(files)):
|
|
|
|
fileName = os.path.splitext(file.split(os.path.sep)[-1])[0] # Load audio and video
|
|
audio, _ = sf.read(os.path.join(video_args.pycropPath, fileName + '.wav'), dtype='float32')
|
|
|
|
video = cv2.VideoCapture(os.path.join(video_args.pycropPath, fileName + '.avi'))
|
|
videoFeature = []
|
|
while video.isOpened():
|
|
ret, frames = video.read()
|
|
if ret == True:
|
|
face = cv2.cvtColor(frames, cv2.COLOR_BGR2GRAY)
|
|
face = cv2.resize(face, (224,224))
|
|
face = face[int(112-(112/2)):int(112+(112/2)), int(112-(112/2)):int(112+(112/2))]
|
|
videoFeature.append(face)
|
|
else:
|
|
break
|
|
|
|
video.release()
|
|
visual = np.array(videoFeature)/255.0
|
|
visual = (visual - 0.4161)/0.1688
|
|
|
|
length = int(audio.shape[0]/16000*25)
|
|
if visual.shape[0] < length:
|
|
visual = np.pad(visual, ((0,int(length - visual.shape[0])),(0,0),(0,0)), mode = 'edge')
|
|
|
|
audio /= np.max(np.abs(audio))
|
|
audio = np.expand_dims(audio, axis=0)
|
|
visual = np.expand_dims(visual, axis=0)
|
|
|
|
inputs = (audio, visual)
|
|
est_source = decode_one_audio_AV_MossFormer2_TSE_16K(video_args.model, inputs, args)
|
|
|
|
est_sources.append(est_source)
|
|
|
|
return est_sources
|
|
|
|
def visualization(tracks, est_sources, video_args):
|
|
# CPU: visulize the result for video format
|
|
flist = glob.glob(os.path.join(video_args.pyframesPath, '*.jpg'))
|
|
flist.sort()
|
|
|
|
|
|
for idx, audio in enumerate(est_sources):
|
|
max_value = np.max(np.abs(audio))
|
|
if max_value >1:
|
|
audio /= max_value
|
|
sf.write(video_args.pycropPath +'/est_%s.wav' %idx, audio, 16000)
|
|
|
|
for tidx, track in enumerate(tracks):
|
|
faces = [[] for i in range(len(flist))]
|
|
for fidx, frame in enumerate(track['track']['frame'].tolist()):
|
|
faces[frame].append({'track':tidx, 's':track['proc_track']['s'][fidx], 'x':track['proc_track']['x'][fidx], 'y':track['proc_track']['y'][fidx]})
|
|
|
|
firstImage = cv2.imread(flist[0])
|
|
fw = firstImage.shape[1]
|
|
fh = firstImage.shape[0]
|
|
vOut = cv2.VideoWriter(os.path.join(video_args.pyaviPath, 'video_only.avi'), cv2.VideoWriter_fourcc(*'XVID'), 25, (fw,fh))
|
|
for fidx, fname in tqdm.tqdm(enumerate(flist), total = len(flist)):
|
|
image = cv2.imread(fname)
|
|
for face in faces[fidx]:
|
|
cv2.rectangle(image, (int(face['x']-face['s']), int(face['y']-face['s'])), (int(face['x']+face['s']), int(face['y']+face['s'])),(0,255,0),10)
|
|
vOut.write(image)
|
|
vOut.release()
|
|
|
|
command = ("ffmpeg -y -i %s -i %s -threads %d -c:v copy -c:a copy %s -loglevel panic" % \
|
|
(os.path.join(video_args.pyaviPath, 'video_only.avi'), (video_args.pycropPath +'/est_%s.wav' %tidx), \
|
|
video_args.nDataLoaderThread, os.path.join(video_args.pyaviPath,'video_out_%s.avi'%tidx)))
|
|
output = subprocess.call(command, shell=True, stdout=None)
|
|
|
|
|
|
|
|
|
|
command = "ffmpeg -i %s %s ;" % (
|
|
os.path.join(video_args.pyaviPath, 'video_out_%s.avi' % tidx),
|
|
os.path.join(video_args.pyaviPath, 'video_est_%s.mp4' % tidx)
|
|
)
|
|
command += f"rm {os.path.join(video_args.pyaviPath, 'video_out_%s.avi' % tidx)}"
|
|
output = subprocess.call(command, shell=True, stdout=None)
|
|
|
|
|
|
command = "ffmpeg -i %s %s ;" % (
|
|
os.path.join(video_args.pyaviPath, 'video.avi'),
|
|
os.path.join(video_args.pyaviPath, 'video_orig.mp4')
|
|
)
|
|
command += f"rm {os.path.join(video_args.pyaviPath, 'video_only.avi')} ;"
|
|
command += f"rm {os.path.join(video_args.pyaviPath, 'video.avi')} ;"
|
|
command += f"rm {os.path.join(video_args.pyaviPath, 'audio.wav')} ;"
|
|
output = subprocess.call(command, shell=True, stdout=None)
|