Spaces:
Running
Running
Delete inference.py
Browse files- inference.py +0 -145
inference.py
DELETED
@@ -1,145 +0,0 @@
|
|
1 |
-
from glob import glob
|
2 |
-
import shutil
|
3 |
-
import torch
|
4 |
-
from time import strftime
|
5 |
-
import os, sys, time
|
6 |
-
from argparse import ArgumentParser
|
7 |
-
|
8 |
-
from src.utils.preprocess import CropAndExtract
|
9 |
-
from src.test_audio2coeff import Audio2Coeff
|
10 |
-
from src.facerender.animate import AnimateFromCoeff
|
11 |
-
from src.generate_batch import get_data
|
12 |
-
from src.generate_facerender_batch import get_facerender_data
|
13 |
-
from src.utils.init_path import init_path
|
14 |
-
|
15 |
-
def main(args):
|
16 |
-
#torch.backends.cudnn.enabled = False
|
17 |
-
|
18 |
-
pic_path = args.source_image
|
19 |
-
audio_path = args.driven_audio
|
20 |
-
save_dir = os.path.join(args.result_dir, strftime("%Y_%m_%d_%H.%M.%S"))
|
21 |
-
os.makedirs(save_dir, exist_ok=True)
|
22 |
-
pose_style = args.pose_style
|
23 |
-
device = args.device
|
24 |
-
batch_size = args.batch_size
|
25 |
-
input_yaw_list = args.input_yaw
|
26 |
-
input_pitch_list = args.input_pitch
|
27 |
-
input_roll_list = args.input_roll
|
28 |
-
ref_eyeblink = args.ref_eyeblink
|
29 |
-
ref_pose = args.ref_pose
|
30 |
-
|
31 |
-
current_root_path = os.path.split(sys.argv[0])[0]
|
32 |
-
|
33 |
-
sadtalker_paths = init_path(args.checkpoint_dir, os.path.join(current_root_path, 'src/config'), args.size, args.old_version, args.preprocess)
|
34 |
-
|
35 |
-
#init model
|
36 |
-
preprocess_model = CropAndExtract(sadtalker_paths, device)
|
37 |
-
|
38 |
-
audio_to_coeff = Audio2Coeff(sadtalker_paths, device)
|
39 |
-
|
40 |
-
animate_from_coeff = AnimateFromCoeff(sadtalker_paths, device)
|
41 |
-
|
42 |
-
#crop image and extract 3dmm from image
|
43 |
-
first_frame_dir = os.path.join(save_dir, 'first_frame_dir')
|
44 |
-
os.makedirs(first_frame_dir, exist_ok=True)
|
45 |
-
print('3DMM Extraction for source image')
|
46 |
-
first_coeff_path, crop_pic_path, crop_info = preprocess_model.generate(pic_path, first_frame_dir, args.preprocess,\
|
47 |
-
source_image_flag=True, pic_size=args.size)
|
48 |
-
if first_coeff_path is None:
|
49 |
-
print("Can't get the coeffs of the input")
|
50 |
-
return
|
51 |
-
|
52 |
-
if ref_eyeblink is not None:
|
53 |
-
ref_eyeblink_videoname = os.path.splitext(os.path.split(ref_eyeblink)[-1])[0]
|
54 |
-
ref_eyeblink_frame_dir = os.path.join(save_dir, ref_eyeblink_videoname)
|
55 |
-
os.makedirs(ref_eyeblink_frame_dir, exist_ok=True)
|
56 |
-
print('3DMM Extraction for the reference video providing eye blinking')
|
57 |
-
ref_eyeblink_coeff_path, _, _ = preprocess_model.generate(ref_eyeblink, ref_eyeblink_frame_dir, args.preprocess, source_image_flag=False)
|
58 |
-
else:
|
59 |
-
ref_eyeblink_coeff_path=None
|
60 |
-
|
61 |
-
if ref_pose is not None:
|
62 |
-
if ref_pose == ref_eyeblink:
|
63 |
-
ref_pose_coeff_path = ref_eyeblink_coeff_path
|
64 |
-
else:
|
65 |
-
ref_pose_videoname = os.path.splitext(os.path.split(ref_pose)[-1])[0]
|
66 |
-
ref_pose_frame_dir = os.path.join(save_dir, ref_pose_videoname)
|
67 |
-
os.makedirs(ref_pose_frame_dir, exist_ok=True)
|
68 |
-
print('3DMM Extraction for the reference video providing pose')
|
69 |
-
ref_pose_coeff_path, _, _ = preprocess_model.generate(ref_pose, ref_pose_frame_dir, args.preprocess, source_image_flag=False)
|
70 |
-
else:
|
71 |
-
ref_pose_coeff_path=None
|
72 |
-
|
73 |
-
#audio2ceoff
|
74 |
-
batch = get_data(first_coeff_path, audio_path, device, ref_eyeblink_coeff_path, still=args.still)
|
75 |
-
coeff_path = audio_to_coeff.generate(batch, save_dir, pose_style, ref_pose_coeff_path)
|
76 |
-
|
77 |
-
# 3dface render
|
78 |
-
if args.face3dvis:
|
79 |
-
from src.face3d.visualize import gen_composed_video
|
80 |
-
gen_composed_video(args, device, first_coeff_path, coeff_path, audio_path, os.path.join(save_dir, '3dface.mp4'))
|
81 |
-
|
82 |
-
#coeff2video
|
83 |
-
data = get_facerender_data(coeff_path, crop_pic_path, first_coeff_path, audio_path,
|
84 |
-
batch_size, input_yaw_list, input_pitch_list, input_roll_list,
|
85 |
-
expression_scale=args.expression_scale, still_mode=args.still, preprocess=args.preprocess, size=args.size)
|
86 |
-
|
87 |
-
result = animate_from_coeff.generate(data, save_dir, pic_path, crop_info, \
|
88 |
-
enhancer=args.enhancer, background_enhancer=args.background_enhancer, preprocess=args.preprocess, img_size=args.size)
|
89 |
-
|
90 |
-
shutil.move(result, save_dir+'.mp4')
|
91 |
-
print('The generated video is named:', save_dir+'.mp4')
|
92 |
-
|
93 |
-
if not args.verbose:
|
94 |
-
shutil.rmtree(save_dir)
|
95 |
-
|
96 |
-
|
97 |
-
if __name__ == '__main__':
|
98 |
-
|
99 |
-
parser = ArgumentParser()
|
100 |
-
parser.add_argument("--driven_audio", default='./examples/driven_audio/bus_chinese.wav', help="path to driven audio")
|
101 |
-
parser.add_argument("--source_image", default='./examples/source_image/full_body_1.png', help="path to source image")
|
102 |
-
parser.add_argument("--ref_eyeblink", default=None, help="path to reference video providing eye blinking")
|
103 |
-
parser.add_argument("--ref_pose", default=None, help="path to reference video providing pose")
|
104 |
-
parser.add_argument("--checkpoint_dir", default='./checkpoints', help="path to output")
|
105 |
-
parser.add_argument("--result_dir", default='./results', help="path to output")
|
106 |
-
parser.add_argument("--pose_style", type=int, default=0, help="input pose style from [0, 46)")
|
107 |
-
parser.add_argument("--batch_size", type=int, default=2, help="the batch size of facerender")
|
108 |
-
parser.add_argument("--size", type=int, default=256, help="the image size of the facerender")
|
109 |
-
parser.add_argument("--expression_scale", type=float, default=1., help="the batch size of facerender")
|
110 |
-
parser.add_argument('--input_yaw', nargs='+', type=int, default=None, help="the input yaw degree of the user ")
|
111 |
-
parser.add_argument('--input_pitch', nargs='+', type=int, default=None, help="the input pitch degree of the user")
|
112 |
-
parser.add_argument('--input_roll', nargs='+', type=int, default=None, help="the input roll degree of the user")
|
113 |
-
parser.add_argument('--enhancer', type=str, default=None, help="Face enhancer, [gfpgan, RestoreFormer]")
|
114 |
-
parser.add_argument('--background_enhancer', type=str, default=None, help="background enhancer, [realesrgan]")
|
115 |
-
parser.add_argument("--cpu", dest="cpu", action="store_true")
|
116 |
-
parser.add_argument("--face3dvis", action="store_true", help="generate 3d face and 3d landmarks")
|
117 |
-
parser.add_argument("--still", action="store_true", help="can crop back to the original videos for the full body aniamtion")
|
118 |
-
parser.add_argument("--preprocess", default='crop', choices=['crop', 'extcrop', 'resize', 'full', 'extfull'], help="how to preprocess the images" )
|
119 |
-
parser.add_argument("--verbose",action="store_true", help="saving the intermedia output or not" )
|
120 |
-
parser.add_argument("--old_version",action="store_true", help="use the pth other than safetensor version" )
|
121 |
-
|
122 |
-
|
123 |
-
# net structure and parameters
|
124 |
-
parser.add_argument('--net_recon', type=str, default='resnet50', choices=['resnet18', 'resnet34', 'resnet50'], help='useless')
|
125 |
-
parser.add_argument('--init_path', type=str, default=None, help='Useless')
|
126 |
-
parser.add_argument('--use_last_fc',default=False, help='zero initialize the last fc')
|
127 |
-
parser.add_argument('--bfm_folder', type=str, default='./checkpoints/BFM_Fitting/')
|
128 |
-
parser.add_argument('--bfm_model', type=str, default='BFM_model_front.mat', help='bfm model')
|
129 |
-
|
130 |
-
# default renderer parameters
|
131 |
-
parser.add_argument('--focal', type=float, default=1015.)
|
132 |
-
parser.add_argument('--center', type=float, default=112.)
|
133 |
-
parser.add_argument('--camera_d', type=float, default=10.)
|
134 |
-
parser.add_argument('--z_near', type=float, default=5.)
|
135 |
-
parser.add_argument('--z_far', type=float, default=15.)
|
136 |
-
|
137 |
-
args = parser.parse_args()
|
138 |
-
|
139 |
-
if torch.cuda.is_available() and not args.cpu:
|
140 |
-
args.device = "cuda"
|
141 |
-
else:
|
142 |
-
args.device = "cpu"
|
143 |
-
|
144 |
-
main(args)
|
145 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|