Spaces:
Running
on
T4
Running
on
T4
""" | |
# Copyright 2020 Adobe | |
# All Rights Reserved. | |
# NOTICE: Adobe permits you to use, modify, and distribute this file in | |
# accordance with the terms of the Adobe license agreement accompanying | |
# it. | |
""" | |
import numpy as np | |
import os | |
import matplotlib.pyplot as plt | |
import cv2 | |
import ffmpeg | |
OTHER_SPECIFIC_VOICE = None | |
class Vis(): | |
def __init__(self, fls, filename, audio_filenam=None, fps=100, frames=1000): | |
# from scipy.signal import savgol_filter | |
# fls = savgol_filter(fls, 21, 3, axis=0) | |
# adj nose | |
# fls[:, 27 * 3:28 * 3] = fls[:, 28 * 3:29 * 3] * 2 - fls[:, 29 * 3:30 * 3] | |
# fls[:, 28 * 3:29 * 3] = fls[:, 27 * 3:28 * 3]*0.75 + fls[:, 31 * 3:32 * 3]*0.25 | |
# fls[:, 29 * 3:30 * 3] = fls[:, 27 * 3:28 * 3]*0.5 + fls[:, 31 * 3:32 * 3]*0.5 | |
# fls[:, 30 * 3:31 * 3] = fls[:, 27 * 3:28 * 3] * 0.25 + fls[:, 31 * 3:32 * 3] * 0.75 | |
fls = fls * 120 | |
fls[:, 0::3] += 200 | |
fls[:, 1::3] += 100 | |
fls = fls.reshape((-1, 68, 3)) | |
fls = fls.astype(int) | |
writer = cv2.VideoWriter(os.path.join('MakeItTalk/examples', 'tmp.mp4'), | |
cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), fps, (400, 400)) | |
frames = np.min((fls.shape[0], frames)) | |
for i in range(frames): #fls.shape[0]): | |
# print(i, fls.shape[0]) | |
frame = np.ones((400, 400, 3), np.uint8) * 0 | |
frame = self.__vis_landmark_on_img__(frame, fls[i]) | |
writer.write(frame) | |
writer.release() | |
if(audio_filenam is not None): | |
print(audio_filenam) | |
os.system('ffmpeg -y -i {} -i {} -strict -2 -shortest {}'.format( | |
os.path.join('MakeItTalk/examples', 'tmp.mp4'), | |
audio_filenam, | |
os.path.join('MakeItTalk/examples', '{}_av.mp4'.format(filename)) | |
)) | |
else: | |
os.system('ffmpeg -y -i {} {}'.format( | |
os.path.join('MakeItTalk/examples', 'tmp.mp4'), | |
os.path.join('MakeItTalk/examples', '{}_av.mp4'.format(filename)) | |
)) | |
os.remove(os.path.join('MakeItTalk/examples', 'tmp.mp4')) | |
def __vis_landmark_on_img__(self, img, shape, linewidth=2): | |
''' | |
Visualize landmark on images. | |
''' | |
def draw_curve(idx_list, color=(0, 255, 0), loop=False, lineWidth=linewidth): | |
for i in idx_list: | |
cv2.line(img, (shape[i, 0], shape[i, 1]), (shape[i + 1, 0], shape[i + 1, 1]), color, lineWidth) | |
if (loop): | |
cv2.line(img, (shape[idx_list[0], 0], shape[idx_list[0], 1]), | |
(shape[idx_list[-1] + 1, 0], shape[idx_list[-1] + 1, 1]), color, lineWidth) | |
# draw_curve(list(range(0, 16)), color=(0, 255, 0)) # jaw | |
# draw_curve(list(range(17, 21)), color=(0, 127, 255)) # eye brow | |
# draw_curve(list(range(22, 26)), color=(0, 127, 255)) | |
# draw_curve(list(range(27, 35)), color=(255, 0, 0)) # nose | |
# draw_curve(list(range(36, 41)), loop=True, color=(204, 0, 204)) # eyes | |
# draw_curve(list(range(42, 47)), loop=True, color=(204, 0, 204)) | |
# draw_curve(list(range(48, 59)), loop=True, color=(0, 0, 255)) # mouth | |
# draw_curve(list(range(60, 67)), loop=True, color=(0, 0, 255)) | |
# draw_curve(list(range(60, 64)), loop=False, color=(0, 0, 255)) | |
draw_curve(list(range(0, 16)), color=(0, 255, 0)) # jaw | |
draw_curve(list(range(17, 21)), color=(0, 255, 0)) # eye brow | |
draw_curve(list(range(22, 26)), color=(0, 255, 0)) | |
draw_curve(list(range(27, 35)), color=(0, 255, 0)) # nose | |
draw_curve(list(range(36, 41)), loop=True, color=(0, 255, 0)) # eyes | |
draw_curve(list(range(42, 47)), loop=True, color=(0, 255, 0)) | |
draw_curve(list(range(48, 59)), loop=True, color=(0, 255, 255)) # mouth | |
draw_curve(list(range(60, 67)), loop=True, color=(255, 255, 0)) | |
draw_curve(list(range(60, 64)), loop=False, color=(0, 0, 255)) | |
return img | |
class Vis_old(): | |
def __init__(self, run_name, pred_fl_filename, audio_filename, av_name='NAME', fps=100, frames=625, | |
postfix='', root_dir=r'E:\Dataset\TalkingToon\Obama', ifsmooth=True, rand_start=0): | |
print(root_dir) | |
self.src_dir = os.path.join(root_dir, r'nn_result/{}'.format(run_name)) | |
self.std_face = np.loadtxt(r'src/dataset/utils/STD_FACE_LANDMARKS.txt') | |
self.std_face = self.std_face.reshape((-1, 204)) | |
fls = np.loadtxt(os.path.join(self.src_dir, pred_fl_filename)) | |
fls = fls * 120 | |
fls[:, 0::3] += 200 | |
fls[:, 1::3] += 100 | |
fls = fls.reshape((-1, 68, 3)) | |
fls = fls.astype(int) | |
writer = cv2.VideoWriter(os.path.join(self.src_dir, 'tmp.mp4'), | |
cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), fps, (400, 400)) | |
frames = np.min((fls.shape[0], frames)) | |
for i in range(frames): #fls.shape[0]): | |
# print(i, fls.shape[0]) | |
frame = np.ones((400, 400, 3), np.uint8) * 0 | |
frame = self.__vis_landmark_on_img__(frame, fls[i]) | |
writer.write(frame) | |
writer.release() | |
if(os.path.exists(os.path.join(root_dir, 'demo_wav', '{}'.format(audio_filename)))): | |
ain = os.path.join(root_dir, 'demo_wav', '{}'.format(audio_filename)) | |
else: | |
ain = os.path.join(root_dir, 'raw_wav', '{}'.format(audio_filename)) | |
# print(ain) | |
# vin = ffmpeg.input(os.path.join(self.src_dir, 'tmp.mp4')).video | |
# ain = ffmpeg.input(ain).audio | |
# out = ffmpeg.output(vin, ain, os.path.join(self.src_dir, '{}_av.mp4'.format(pred_fl_filename[:-4])), shortest=None) | |
# out = out.overwrite_output().global_args('-loglevel', 'quiet') | |
# out.run() | |
os.system('ffmpeg -y -loglevel error -i {} -ss {} {}'.format( | |
ain, rand_start/62.5, | |
os.path.join(self.src_dir, '{}_a_tmp.wav'.format(av_name)) | |
)) | |
os.system('ffmpeg -y -loglevel error -i {} -i {} -pix_fmt yuv420p -strict -2 -shortest {}'.format( | |
os.path.join(self.src_dir, 'tmp.mp4'), | |
os.path.join(self.src_dir, '{}_a_tmp.wav'.format(av_name)), | |
os.path.join(self.src_dir, '{}_av.mp4'.format(av_name)) | |
)) | |
os.remove(os.path.join(self.src_dir, 'tmp.mp4')) | |
os.remove(os.path.join(self.src_dir, '{}_a_tmp.wav'.format(av_name))) | |
# os.remove(os.path.join(self.src_dir, filename)) | |
# exit(0) | |
def __vis_landmark_on_img__(self, img, shape, linewidth=2): | |
''' | |
Visualize landmark on images. | |
''' | |
def draw_curve(idx_list, color=(0, 255, 0), loop=False, lineWidth=linewidth): | |
for i in idx_list: | |
cv2.line(img, (shape[i, 0], shape[i, 1]), (shape[i + 1, 0], shape[i + 1, 1]), color, lineWidth) | |
if (loop): | |
cv2.line(img, (shape[idx_list[0], 0], shape[idx_list[0], 1]), | |
(shape[idx_list[-1] + 1, 0], shape[idx_list[-1] + 1, 1]), color, lineWidth) | |
# draw_curve(list(range(0, 16)), color=(0, 255, 0)) # jaw | |
# draw_curve(list(range(17, 21)), color=(0, 127, 255)) # eye brow | |
# draw_curve(list(range(22, 26)), color=(0, 127, 255)) | |
# draw_curve(list(range(27, 35)), color=(255, 0, 0)) # nose | |
# draw_curve(list(range(36, 41)), loop=True, color=(204, 0, 204)) # eyes | |
# draw_curve(list(range(42, 47)), loop=True, color=(204, 0, 204)) | |
# draw_curve(list(range(48, 59)), loop=True, color=(0, 0, 255)) # mouth | |
# draw_curve(list(range(60, 67)), loop=True, color=(0, 0, 255)) | |
# draw_curve(list(range(60, 64)), loop=False, color=(0, 0, 255)) | |
draw_curve(list(range(0, 16)), color=(0, 255, 0)) # jaw | |
draw_curve(list(range(17, 21)), color=(0, 255, 0)) # eye brow | |
draw_curve(list(range(22, 26)), color=(0, 255, 0)) | |
draw_curve(list(range(27, 35)), color=(0, 255, 0)) # nose | |
draw_curve(list(range(36, 41)), loop=True, color=(0, 255, 0)) # eyes | |
draw_curve(list(range(42, 47)), loop=True, color=(0, 255, 0)) | |
draw_curve(list(range(48, 59)), loop=True, color=(0, 255, 255)) # mouth | |
draw_curve(list(range(60, 67)), loop=True, color=(255, 255, 0)) | |
draw_curve(list(range(60, 64)), loop=False, color=(0, 0, 255)) | |
return img | |
class Vis_comp(): | |
def __init__(self, run_name, pred1, pred2, audio_filename, av_name='NAME', fps=100, frames=625, postfix='', root_dir=r'E:\Dataset\TalkingToon\Obama', ifsmooth=True): | |
print(root_dir) | |
self.src_dir = os.path.join(root_dir, r'nn_result/{}'.format(run_name)) | |
self.std_face = np.loadtxt(r'src/dataset/utils/STD_FACE_LANDMARKS.txt') | |
self.std_face = self.std_face.reshape((-1, 204)) | |
def fls_adj(fls): | |
fls = fls * 120 | |
fls[:, 0::3] += 200 | |
fls[:, 1::3] += 100 | |
fls = fls.reshape((-1, 68, 3)) | |
fls = fls.astype(int) | |
return fls | |
fls = np.loadtxt(os.path.join(self.src_dir, pred1)) | |
fls2 = np.loadtxt(os.path.join(self.src_dir, pred2)) | |
fls = fls_adj(fls) | |
fls2 = fls_adj(fls2) | |
writer = cv2.VideoWriter(os.path.join(self.src_dir, 'tmp.mp4'), | |
cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), fps, (400, 400)) | |
frames = np.min((fls.shape[0], frames)) | |
for i in range(frames): #fls.shape[0]): | |
# print(i, fls.shape[0]) | |
frame = np.ones((400, 400, 3), np.uint8) * 0 | |
frame = self.__vis_landmark_on_img__(frame, fls[i]) | |
frame = self.__vis_landmark_on_img__(frame, fls2[i]) | |
writer.write(frame) | |
writer.release() | |
if(os.path.exists(os.path.join(root_dir, 'demo_wav', '{}'.format(audio_filename)))): | |
ain = os.path.join(root_dir, 'demo_wav', '{}'.format(audio_filename)) | |
else: | |
ain = os.path.join(root_dir, 'raw_wav', '{}'.format(audio_filename)) | |
os.system('ffmpeg -y -loglevel error -i {} -i {} -pix_fmt yuv420p -strict -2 -shortest {}'.format( | |
os.path.join(self.src_dir, 'tmp.mp4'), | |
ain, | |
os.path.join(self.src_dir, '{}_av.mp4'.format(av_name)) | |
)) | |
os.remove(os.path.join(self.src_dir, 'tmp.mp4')) | |
def __vis_landmark_on_img__(self, img, shape, linewidth=2): | |
''' | |
Visualize landmark on images. | |
''' | |
def draw_curve(idx_list, color=(0, 255, 0), loop=False, lineWidth=linewidth): | |
for i in idx_list: | |
cv2.line(img, (shape[i, 0], shape[i, 1]), (shape[i + 1, 0], shape[i + 1, 1]), color, lineWidth) | |
if (loop): | |
cv2.line(img, (shape[idx_list[0], 0], shape[idx_list[0], 1]), | |
(shape[idx_list[-1] + 1, 0], shape[idx_list[-1] + 1, 1]), color, lineWidth) | |
draw_curve(list(range(0, 16)), color=(0, 255, 0)) # jaw | |
draw_curve(list(range(17, 21)), color=(0, 255, 0)) # eye brow | |
draw_curve(list(range(22, 26)), color=(0, 255, 0)) | |
draw_curve(list(range(27, 35)), color=(0, 255, 0)) # nose | |
draw_curve(list(range(36, 41)), loop=True, color=(0, 255, 0)) # eyes | |
draw_curve(list(range(42, 47)), loop=True, color=(0, 255, 0)) | |
draw_curve(list(range(48, 59)), loop=True, color=(0, 255, 255)) # mouth | |
draw_curve(list(range(60, 67)), loop=True, color=(255, 255, 0)) | |
draw_curve(list(range(60, 64)), loop=False, color=(0, 0, 255)) | |
return img |