-
Notifications
You must be signed in to change notification settings - Fork 0
/
render_final_video.py
115 lines (92 loc) · 3.9 KB
/
render_final_video.py
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
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
import argparse
import json
import os
import re
from PIL import Image
def atoi(text):
return int(text) if text.isdigit() else text
def natural_keys(text):
"""
alist.sort(key=natural_keys) sorts in human order
http://nedbatchelder.com/blog/200712/human_sorting.html
(See Toothy's implementation in the comments)
"""
return [atoi(c) for c in re.split(r"(\d+)", text)]
def render_final_video(output_dir):
for file in os.listdir(output_dir):
if file.endswith(".json"):
render_name = file.split(".")[0]
frame_dir = os.path.join(output_dir, render_name)
os.makedirs(frame_dir, exist_ok=True)
with open(os.path.join(output_dir, file), "r") as f:
data = json.load(f)
speaker_gt_frame = data["sp_gt_frame_address"]
listener_gt_frame = data["lt_gt_frame_address"]
mesh_frame_dir = data["mesh_video_address"]
fake_frame_dir = data["fake_video_address"]
if not os.path.isdir(mesh_frame_dir):
mesh_frame = [
img
for img in os.listdir(os.path.join(output_dir, mesh_frame_dir))
if img.endswith(".png")
]
fake_frame = [
img
for img in os.listdir(os.path.join(output_dir, fake_frame_dir))
if img.endswith(".png")
]
else:
mesh_frame = [
img
for img in os.listdir( mesh_frame_dir)
if img.endswith(".png")
]
fake_frame = [
img
for img in os.listdir(fake_frame_dir)
if img.endswith(".png")
]
mesh_frame.sort(key=natural_keys)
fake_frame.sort(key=natural_keys)
for frame_idx in range(len(mesh_frame)):
if not os.path.isdir(mesh_frame_dir):
mesh_img = os.path.join(
output_dir, mesh_frame_dir, mesh_frame[frame_idx]
)
fake_img = os.path.join(
output_dir, fake_frame_dir, fake_frame[frame_idx]
)
else:
mesh_img = os.path.join(
mesh_frame_dir, mesh_frame[frame_idx]
)
fake_img = os.path.join(
fake_frame_dir, fake_frame[frame_idx]
)
speaker_gt_img = speaker_gt_frame[frame_idx]
listener_gt_img = listener_gt_frame[frame_idx]
# read image use pil
mesh_img = Image.open(mesh_img)
fake_img = Image.open(fake_img)
speaker_gt_img = Image.open(speaker_gt_img)
listener_gt_img = Image.open(listener_gt_img)
# combine image
combined_img = Image.new(
"RGB", (mesh_img.width * 4, mesh_img.height * 1)
)
combined_img.paste(speaker_gt_img, (0, 0))
combined_img.paste(listener_gt_img, (mesh_img.width, 0))
combined_img.paste(mesh_img, (mesh_img.width * 2, 0))
combined_img.paste(fake_img, (mesh_img.width * 3, 0))
# save image
combined_img.save(os.path.join(frame_dir, f"frame_{frame_idx}.png"))
# create video with high bit rate
os.system(
f"ffmpeg -r 30 -i {frame_dir}/frame_%d.png -b:v 10000k -vcodec mpeg4 -y {frame_dir}.mp4"
)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--output-dir", type=str, default="output/")
args = parser.parse_args()
print("RENDERING FINAL VIDEO...")
render_final_video(args.output_dir)