Update inference.py

inference.py

@@ -1,295 +1,280 @@
+from os import listdir, path
 import numpy as np
-import cv2
-import os
-import sys
-import argparse
-import audio
+import scipy, cv2, os, sys, argparse, audio
+import json, subprocess, random, string
 from tqdm import tqdm
 from glob import glob
-import torch
-import face_detection
+import torch, face_detection
 from models import Wav2Lip
 import platform
 
-try:
-    last_face = cv2.imread("last_face.jpg")
-except Exception as e:
-    last_face = None
+parser = argparse.ArgumentParser(description='Inference code to lip-sync videos in the wild using Wav2Lip models')
 
-if last_face is None:
-    last_face = cv2.imread("imagen_por_defecto.jpg")
+parser.add_argument('--checkpoint_path', type=str, 
+					help='Name of saved checkpoint to load weights from', required=True)
 
-parser = argparse.ArgumentParser(description='Inference code to lip-sync videos in the wild using Wav2Lip models')
+parser.add_argument('--face', type=str, 
+					help='Filepath of video/image that contains faces to use', required=True)
+parser.add_argument('--audio', type=str, 
+					help='Filepath of video/audio file to use as raw audio source', required=True)
+parser.add_argument('--outfile', type=str, help='Video path to save result. See default for an e.g.', 
+								default='results/result_voice.mp4')
+
+parser.add_argument('--static', type=bool, 
+					help='If True, then use only first video frame for inference', default=False)
+parser.add_argument('--fps', type=float, help='Can be specified only if input is a static image (default: 25)', 
+					default=25., required=False)
+
+parser.add_argument('--pads', nargs='+', type=int, default=[0, 10, 0, 0], 
+					help='Padding (top, bottom, left, right). Please adjust to include chin at least')
+
+parser.add_argument('--face_det_batch_size', type=int, 
+					help='Batch size for face detection', default=16)
+parser.add_argument('--wav2lip_batch_size', type=int, help='Batch size for Wav2Lip model(s)', default=128)
 
-parser.add_argument('--checkpoint_path', type=str, required=True)
-parser.add_argument('--face', type=str, required=True)
-parser.add_argument('--audio', type=str, required=True)
-parser.add_argument('--outfile', type=str, default='results/result_voice.mp4')
-parser.add_argument('--static', type=bool, default=False)
-parser.add_argument('--fps', type=float, default=25., required=False)
-parser.add_argument('--pads', nargs='+', type=int, default=[0, 10, 0, 0])
-parser.add_argument('--face_det_batch_size', type=int, default=16)
-parser.add_argument('--wav2lip_batch_size', type=int, default=128)
-parser.add_argument('--resize_factor', default=1, type=int)
-parser.add_argument('--crop', nargs='+', type=int, default=[0, -1, 0, -1])
-parser.add_argument('--box', nargs='+', type=int, default=[-1, -1, -1, -1])
-parser.add_argument('--rotate', default=False, action='store_true')
-parser.add_argument('--nosmooth', default=False, action='store_true')
+parser.add_argument('--resize_factor', default=1, type=int, 
+			help='Reduce the resolution by this factor. Sometimes, best results are obtained at 480p or 720p')
+
+parser.add_argument('--crop', nargs='+', type=int, default=[0, -1, 0, -1], 
+					help='Crop video to a smaller region (top, bottom, left, right). Applied after resize_factor and rotate arg. ' 
+					'Useful if multiple face present. -1 implies the value will be auto-inferred based on height, width')
+
+parser.add_argument('--box', nargs='+', type=int, default=[-1, -1, -1, -1], 
+					help='Specify a constant bounding box for the face. Use only as a last resort if the face is not detected.'
+					'Also, might work only if the face is not moving around much. Syntax: (top, bottom, left, right).')
+
+parser.add_argument('--rotate', default=False, action='store_true',
+					help='Sometimes videos taken from a phone can be flipped 90deg. If true, will flip video right by 90deg.'
+					'Use if you get a flipped result, despite feeding a normal looking video')
+
+parser.add_argument('--nosmooth', default=False, action='store_true',
+					help='Prevent smoothing face detections over a short temporal window')
 
 args = parser.parse_args()
 args.img_size = 96
 
 if os.path.isfile(args.face) and args.face.split('.')[1] in ['jpg', 'png', 'jpeg']:
-    args.static = True
+	args.static = True
 
 def get_smoothened_boxes(boxes, T):
-    for i in range(len(boxes)):
-        if i + T > len(boxes):
-            window = boxes[len(boxes) - T:]
-        else:
-            window = boxes[i : i + T]
-        boxes[i] = np.mean(window, axis=0)
-    return boxes
+	for i in range(len(boxes)):
+		if i + T > len(boxes):
+			window = boxes[len(boxes) - T:]
+		else:
+			window = boxes[i : i + T]
+		boxes[i] = np.mean(window, axis=0)
+	return boxes
 
 def face_detect(images):
-    detector = face_detection.FaceAlignment(face_detection.LandmarksType._2D, flip_input=False, device=device)
-    batch_size = args.face_det_batch_size
-    last_face = None
-
-    while 1:
-        predictions = []
-        try:
-            for i in tqdm(range(0, len(images), batch_size)):
-                predictions.extend(detector.get_detections_for_batch(np.array(images[i:i + batch_size])))
-        except RuntimeError:
-            if batch_size == 1:
-                raise RuntimeError('Image too big to run face detection on GPU. Please use the --resize_factor argument')
-            batch_size //= 2
-            continue
-        break
-    head_exist = []
-    results = []
-    pady1, pady2, padx1, padx2 = args.pads
-
-    first_head_rect = None
-    first_head_image = None
-    for rect, image in zip(predictions, images):
-        if rect is not None:
-            first_head_rect = rect
-            first_head_image = image
-            break
-    for rect, image in zip(predictions, images):
-        if rect is None:
-            head_exist.append(False)
-            if len(results) == 0:
-                y1 = max(0, first_head_rect[1] - pady1)
-                y2 = min(first_head_image.shape[0], first_head_rect[3] + pady2)
-                x1 = max(0, first_head_rect[0] - padx1)
-                x2 = min(first_head_image.shape[1], first_head_rect[2] + padx2)
-                results.append([x1, y1, x2, y2])
-            else:
-                results.append(results[-1])
-        else:
-            head_exist.append(True)
-            y1 = max(0, rect[1] - pady1)
-            y2 = min(image.shape[0], rect[3] + pady2)
-            x1 = max(0, rect[0] - padx1)
-            x2 = min(image.shape[1], rect[2] + padx2)
-            results.append([x1, y1, x2, y2])
-            last_face = image[y1: y2, x1:x2]
-            cv2.imwrite("last_face.jpg", last_face)
-
-    boxes = np.array(results)
-    if not args.nosmooth:
-        boxes = get_smoothened_boxes(boxes, T=5)
-    results = [[image[y1: y2, x1:x2], (y1, y2, x1, x2)] for image, (x1, y1, x2, y2) in zip(images, boxes)]
-
-    del detector
-    return results, head_exist
+	detector = face_detection.FaceAlignment(face_detection.LandmarksType._2D, 
+											flip_input=False, device=device)
+
+	batch_size = args.face_det_batch_size
+	
+	while 1:
+		predictions = []
+		try:
+			for i in tqdm(range(0, len(images), batch_size)):
+				predictions.extend(detector.get_detections_for_batch(np.array(images[i:i + batch_size])))
+		except RuntimeError:
+			if batch_size == 1: 
+				raise RuntimeError('Image too big to run face detection on GPU. Please use the --resize_factor argument')
+			batch_size //= 2
+			print('Recovering from OOM error; New batch size: {}'.format(batch_size))
+			continue
+		break
+
+	results = []
+	pady1, pady2, padx1, padx2 = args.pads
+	for rect, image in zip(predictions, images):
+		if rect is None:
+			cv2.imwrite('temp/faulty_frame.jpg', image) # check this frame where the face was not detected.
+			raise ValueError('Face not detected! Ensure the video contains a face in all the frames.')
+
+		y1 = max(0, rect[1] - pady1)
+		y2 = min(image.shape[0], rect[3] + pady2)
+		x1 = max(0, rect[0] - padx1)
+		x2 = min(image.shape[1], rect[2] + padx2)
+		
+		results.append([x1, y1, x2, y2])
+
+	boxes = np.array(results)
+	if not args.nosmooth: boxes = get_smoothened_boxes(boxes, T=5)
+	results = [[image[y1: y2, x1:x2], (y1, y2, x1, x2)] for image, (x1, y1, x2, y2) in zip(images, boxes)]
+
+	del detector
+	return results 
 
 def datagen(frames, mels):
-    img_batch, head_exist_batch, mel_batch, frame_batch, coords_batch = [], [], [], [], []
-
-    if args.box[0] == -1:
-        if not args.static:
-            face_det_results, head_exist = face_detect(frames)
-        else:
-            face_det_results, head_exist = face_detect([frames[0]])
-    else:
-        y1, y2, x1, x2 = args.box
-        face_det_results = [[f[y1: y2, x1:x2], (y1, y2, x1, x2)] for f in frames]
-        head_exist = [True] * len(frames)
-
-    for i, m in enumerate(mels):
-        idx = 0 if args.static else i % len(frames)
-        frame_to_save = frames[idx].copy()
-        face, coords = face_det_results[idx].copy()
-
-        face = cv2.resize(face, (args.img_size, args.img_size))
-        head_exist_batch.append(head_exist[idx])
-        img_batch.append(face)
-        melspec = m
-        mel_batch.append(melspec)
-        frame_batch.append(frame_to_save)
-        coords_batch.append(coords)
-
-        if len(img_batch) >= args.wav2lip_batch_size:
-            img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)
-
-            img_masked = img_batch.copy()
-            img_masked[:, args.img_size // 2:] = 0
-
-            img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
-            mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])
-
-            yield img_batch, head_exist_batch, mel_batch, frame_batch, coords_batch
-            img_batch, head_exist_batch, mel_batch, frame_batch, coords_batch = [], [], [], [], []
-
-    last_face = cv2.imread("last_face.jpg")
-    last_face = cv2.resize(last_face, (args.img_size, args.img_size))
-    img_batch.append(last_face)
-    melspec = mels[-1]
-    mel_batch.append(melspec)
-    frame_batch.append(frames[-1])
-    coords_batch.append(face_det_results[-1][1])
-    head_exist_batch.append(head_exist[-1])
-
-    if len(img_batch) > 0:
-        img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)
-
-        img_masked = img_batch.copy()
-        img_masked[:, args.img_size // 2:] = 0
-
-        img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
-        mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])
-
-        yield img_batch, head_exist_batch, mel_batch, frame_batch, coords_batch
+	img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []
+
+	if args.box[0] == -1:
+		if not args.static:
+			face_det_results = face_detect(frames) # BGR2RGB for CNN face detection
+		else:
+			face_det_results = face_detect([frames[0]])
+	else:
+		print('Using the specified bounding box instead of face detection...')
+		y1, y2, x1, x2 = args.box
+		face_det_results = [[f[y1: y2, x1:x2], (y1, y2, x1, x2)] for f in frames]
+
+	for i, m in enumerate(mels):
+		idx = 0 if args.static else i%len(frames)
+		frame_to_save = frames[idx].copy()
+		face, coords = face_det_results[idx].copy()
+
+		face = cv2.resize(face, (args.img_size, args.img_size))
+			
+		img_batch.append(face)
+		mel_batch.append(m)
+		frame_batch.append(frame_to_save)
+		coords_batch.append(coords)
+
+		if len(img_batch) >= args.wav2lip_batch_size:
+			img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)
+
+			img_masked = img_batch.copy()
+			img_masked[:, args.img_size//2:] = 0
+
+			img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
+			mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])
+
+			yield img_batch, mel_batch, frame_batch, coords_batch
+			img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []
+
+	if len(img_batch) > 0:
+		img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)
+
+		img_masked = img_batch.copy()
+		img_masked[:, args.img_size//2:] = 0
+
+		img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
+		mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])
+
+		yield img_batch, mel_batch, frame_batch, coords_batch
 
 mel_step_size = 16
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
+print('Using {} for inference.'.format(device))
 
 def _load(checkpoint_path):
-    if device == 'cuda':
-        checkpoint = torch.load(checkpoint_path)
-    else:
-        checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage)
-    return checkpoint
+	if device == 'cuda':
+		checkpoint = torch.load(checkpoint_path)
+	else:
+		checkpoint = torch.load(checkpoint_path,
+								map_location=lambda storage, loc: storage)
+	return checkpoint
 
 def load_model(path):
-    model = Wav2Lip()
-    print("Load checkpoint from: {}".format(path))
-    checkpoint = _load(path)
-    s = checkpoint["state_dict"]
-    new_s = {}
-    for k, v in s.items():
-        new_s[k.replace('module.', '')] = v
-    model.load_state_dict(new_s)
-    model = model.to(device)
-    return model.eval()
+	model = Wav2Lip()
+	print("Load checkpoint from: {}".format(path))
+	checkpoint = _load(path)
+	s = checkpoint["state_dict"]
+	new_s = {}
+	for k, v in s.items():
+		new_s[k.replace('module.', '')] = v
+	model.load_state_dict(new_s)
+
+	model = model.to(device)
+	return model.eval()
 
 def main():
-    if not os.path.isfile(args.face):
-        raise ValueError('--face argument must be a valid path to video/image file')
+	if not os.path.isfile(args.face):
+		raise ValueError('--face argument must be a valid path to video/image file')
 
-    elif args.face.split('.')[1] in ['jpg', 'png', 'jpeg']:
-        full_frames = [cv2.imread(args.face)]
-        fps = args.fps
+	elif args.face.split('.')[1] in ['jpg', 'png', 'jpeg']:
+		full_frames = [cv2.imread(args.face)]
+		fps = args.fps
 
-    else:
-        video_stream = cv2.VideoCapture(args.face)
-        fps = video_stream.get(cv2.CAP_PROP_FPS)
+	else:
+		video_stream = cv2.VideoCapture(args.face)
+		fps = video_stream.get(cv2.CAP_PROP_FPS)
 
-        print('Reading video frames...')
+		print('Reading video frames...')
 
-        full_frames = []
-        while 1:
-            still_reading, frame = video_stream.read()
-            if not still_reading:
-                video_stream.release()
-                break
-            if args.resize_factor > 1:
-                frame = cv2.resize(frame, (frame.shape[1] // args.resize_factor, frame.shape[0] // args.resize_factor))
+		full_frames = []
+		while 1:
+			still_reading, frame = video_stream.read()
+			if not still_reading:
+				video_stream.release()
+				break
+			if args.resize_factor > 1:
+				frame = cv2.resize(frame, (frame.shape[1]//args.resize_factor, frame.shape[0]//args.resize_factor))
 
-            if args.rotate:
-                frame = cv2.rotate(frame, cv2.cv2.ROTATE_90_CLOCKWISE)
+			if args.rotate:
+				frame = cv2.rotate(frame, cv2.cv2.ROTATE_90_CLOCKWISE)
 
-            y1, y2, x1, x2 = args.crop
-            if x2 == -1: x2 = frame.shape[1]
-            if y2 == -1: y2 = frame.shape[0]
+			y1, y2, x1, x2 = args.crop
+			if x2 == -1: x2 = frame.shape[1]
+			if y2 == -1: y2 = frame.shape[0]
 
-            frame = frame[y1:y2, x1:x2]
+			frame = frame[y1:y2, x1:x2]
 
-            full_frames.append(frame)
+			full_frames.append(frame)
 
-    print ("Number of frames available for inference: "+str(len(full_frames)))
+	print ("Number of frames available for inference: "+str(len(full_frames)))
 
-    if not args.audio.endswith('.wav'):
-        print('Extracting raw audio...')
-        command = 'ffmpeg -y -i {} -strict -2 {}'.format(args.audio, 'temp/temp.wav')
+	if not args.audio.endswith('.wav'):
+		print('Extracting raw audio...')
+		command = 'ffmpeg -y -i {} -strict -2 {}'.format(args.audio, 'temp/temp.wav')
 
-        subprocess.call(command, shell=True)
-        args.audio = 'temp/temp.wav'
+		subprocess.call(command, shell=True)
+		args.audio = 'temp/temp.wav'
 
-    wav = audio.load_wav(args.audio, 16000)
-    mel = audio.melspectrogram(wav)
-    print(mel.shape)
+	wav = audio.load_wav(args.audio, 16000)
+	mel = audio.melspectrogram(wav)
+	print(mel.shape)
 
-    if np.isnan(mel.reshape(-1)).sum() > 0:
-        raise ValueError('Mel contains nan! Using a TTS voice? Add a small epsilon noise to the wav file and try again')
+	if np.isnan(mel.reshape(-1)).sum() > 0:
+		raise ValueError('Mel contains nan! Using a TTS voice? Add a small epsilon noise to the wav file and try again')
 
-    mel_chunks = []
-    mel_idx_multiplier = 80. / fps
-    i = 0
-    while 1:
-        start_idx = int(i * mel_idx_multiplier)
-        if start_idx + mel_step_size > len(mel[0]):
-            mel_chunks.append(mel[:, len(mel[0]) - mel_step_size:])
-            break
-        mel_chunks.append(mel[:, start_idx : start_idx + mel_step_size])
-        i += 1
+	mel_chunks = []
+	mel_idx_multiplier = 80./fps 
+	i = 0
+	while 1:
+		start_idx = int(i * mel_idx_multiplier)
+		if start_idx + mel_step_size > len(mel[0]):
+			mel_chunks.append(mel[:, len(mel[0]) - mel_step_size:])
+			break
+		mel_chunks.append(mel[:, start_idx : start_idx + mel_step_size])
+		i += 1
 
-    print("Length of mel chunks: {}".format(len(mel_chunks)))
+	print("Length of mel chunks: {}".format(len(mel_chunks)))
 
-    full_frames = full_frames[:len(mel_chunks)]
+	full_frames = full_frames[:len(mel_chunks)]
 
-    batch_size = args.wav2lip_batch_size
-    gen = datagen(full_frames.copy(), mel_chunks)
+	batch_size = args.wav2lip_batch_size
+	gen = datagen(full_frames.copy(), mel_chunks)
 
-    for i, (img_batch, exist_head_batch, mel_batch, frames, coords) in enumerate(tqdm(gen,
-                                            total=int(np.ceil(float(len(mel_chunks))/batch_size)))):
-        if i == 0:
-            model = load_model(args.checkpoint_path)
-            print("Model loaded")
+	for i, (img_batch, mel_batch, frames, coords) in enumerate(tqdm(gen, 
+											total=int(np.ceil(float(len(mel_chunks))/batch_size)))):
+		if i == 0:
+			model = load_model(args.checkpoint_path)
+			print ("Model loaded")
 
-            frame_h, frame_w = full_frames[0].shape[:-1]
-            out = cv2.VideoWriter('temp/result.avi', 
-                                    cv2.VideoWriter_fourcc(*'DIVX'), fps, (frame_w, frame_h))
+			frame_h, frame_w = full_frames[0].shape[:-1]
+			out = cv2.VideoWriter('temp/result.avi', 
+									cv2.VideoWriter_fourcc(*'DIVX'), fps, (frame_w, frame_h))
 
-        img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device)
-        mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to(device)
+		img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device)
+		mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to(device)
 
-        with torch.no_grad():
-            pred = model(mel_batch, img_batch)
+		with torch.no_grad():
+			pred = model(mel_batch, img_batch)
 
-        pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255.
+		pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255.
+		
+		for p, f, c in zip(pred, frames, coords):
+			y1, y2, x1, x2 = c
+			p = cv2.resize(p.astype(np.uint8), (x2 - x1, y2 - y1))
 
-        i = 0
-        for p, f, c, exist in zip(pred, frames, coords, exist_head_batch):
-            i += 1
-            if not exist:
-                out.write(f)
-            else:
-                y1, y2, x1, x2 = c
-                p = cv2.resize(p.astype(np.uint8), (x2 - x1, y2 - y1))
-                head_high, head_width, _ = p.shape
-                width_cut = int(head_width * 0.2)
-                f[y1:y2, x1+width_cut:x2-width_cut] = p[:, width_cut:head_width-width_cut]
-                out.write(f)
+			f[y1:y2, x1:x2] = p
+			out.write(f)
 
-    out.release()
+	out.release()
 
-    command = 'ffmpeg -y -i {} -i {} -strict -2 -q:v 1 {}'.format(args.audio, 'temp/result.avi', args.outfile)
-    subprocess.call(command, shell=platform.system() != 'Windows')
+	command = 'ffmpeg -y -i {} -i {} -strict -2 -q:v 1 {}'.format(args.audio, 'temp/result.avi', args.outfile)
+	subprocess.call(command, shell=platform.system() != 'Windows')
 
 if __name__ == '__main__':
-    main()
+	main()