File size: 4,177 Bytes
c244e32
712165a
 
c244e32
712165a
c244e32
 
23e13f8
712165a
23e13f8
c244e32
 
 
 
 
 
 
 
 
 
 
712165a
23e13f8
c244e32
 
23e13f8
c244e32
 
 
 
712165a
c244e32
 
 
 
712165a
c244e32
23e13f8
712165a
 
 
 
 
 
 
 
 
c244e32
712165a
c244e32
712165a
c244e32
712165a
 
c244e32
712165a
 
 
c244e32
23e13f8
c244e32
 
712165a
c244e32
712165a
 
 
 
 
 
c244e32
712165a
c244e32
712165a
 
 
 
c244e32
712165a
 
c244e32
712165a
c244e32
 
712165a
 
c244e32
712165a
c244e32
 
23e13f8
712165a
 
 
 
c244e32
712165a
 
 
 
c244e32
712165a
c244e32
23e13f8
712165a
 
 
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
import os
import cv2
import numpy as np
import tempfile
from tqdm import tqdm
from subprocess import call

import trimesh
from scipy.io import wavfile
import librosa

class Struct(object):
    def __init__(self, **kwargs):
        for key, val in kwargs.items():
            setattr(self, key, val)

def add_image_text(img, text, color=(0,0,255), w=800, h=800):
    font = cv2.FONT_HERSHEY_SIMPLEX
    img = np.require(img, dtype='f4', requirements=['O', 'W'])
    img.flags.writeable = True
    img1 = img.copy()
    img1 = cv2.putText(img1, '%s' % (text), (50, 100), font, 2, color, 2, 1)
    img1 = cv2.rectangle(img1, (0, 0), (w, h), color, thickness=3)
    return img1

class RenderTool:
    def __init__(self, out_path):
        path = os.path.join(os.getcwd(), 'visualise/smplx/SMPLX_NEUTRAL.npz')
        model_data = np.load(path, allow_pickle=True)
        data_struct = Struct(**model_data)
        self.f = data_struct.f  # faces
        self.out_path = out_path
        if not os.path.exists(self.out_path):
            os.makedirs(self.out_path)

    def _render_sequences(self, cur_wav_file, v_list, j=-1, run_in_parallel=False):
        symbol = '/'
        print("Render {} {} sequence.".format(cur_wav_file.split(symbol)[-2], cur_wav_file.split(symbol)[-1]))
        
        directory = os.path.join(self.out_path, cur_wav_file.split(symbol)[2].split(symbol)[0])
        if not os.path.exists(directory):
            os.makedirs(directory)
        
        if j == -1:
            video_fname = os.path.join(directory, '%s.mp4' % cur_wav_file.split(symbol)[-1].split('.')[-2].split(symbol)[-1])
        elif j == -2:
            video_fname = os.path.join(directory, cur_wav_file.split(symbol)[-3]+'--%s.mp4' % cur_wav_file.split(symbol)[-1].split('.')[-2].split(symbol)[-1])
        else:
            video_fname = os.path.join(directory, str(j)+'_%s.mp4' % cur_wav_file.split(symbol)[-1].split('.')[-2].split(symbol)[-1])

        self._render_sequences_helper(video_fname, cur_wav_file, v_list)

    def _render_sequences_helper(self, video_fname, cur_wav_file, v_list):
        num_frames = v_list[0].shape[0]

        # Prepare output frames folder
        frames_dir = os.path.join(os.path.dirname(video_fname), "frames_tmp")
        os.makedirs(frames_dir, exist_ok=True)

        center = np.mean(v_list[0][0], axis=0)

        for i_frame in tqdm(range(num_frames)):
            cur_img_list = []
            for i in range(len(v_list)):
                mesh = trimesh.Trimesh(vertices=v_list[i][i_frame], faces=self.f)
                scene = trimesh.Scene(mesh)
                png = scene.save_image(resolution=[800, 800], visible=True)
                
                img_array = np.asarray(bytearray(png), dtype=np.uint8)
                img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)

                cur_img_list.append(img)

            if len(cur_img_list) == 1:
                final_img = cur_img_list[0]
            else:
                final_img = np.hstack(cur_img_list)

            frame_path = os.path.join(frames_dir, f"frame_{i_frame:04d}.png")
            cv2.imwrite(frame_path, final_img)

        # Create video from frames
        tmp_audio_file = tempfile.NamedTemporaryFile('w', suffix='.wav', dir=os.path.dirname(video_fname))
        tmp_audio_file.close()

        audio, sr = librosa.load(cur_wav_file, sr=16000)
        wavfile.write(tmp_audio_file.name, sr, audio)

        tmp_video_file = tempfile.NamedTemporaryFile('w', suffix='.mp4', dir=os.path.dirname(video_fname))
        tmp_video_file.close()

        cmd_frames_to_video = (
            f"ffmpeg -framerate 30 -i {frames_dir}/frame_%04d.png -c:v libx264 -pix_fmt yuv420p {tmp_video_file.name}"
        ).split()
        call(cmd_frames_to_video)

        cmd_merge_audio = (
            f"ffmpeg -i {tmp_video_file.name} -i {tmp_audio_file.name} -vcodec copy -acodec aac {video_fname}"
        ).split()
        call(cmd_merge_audio)

        # Cleanup
        os.remove(tmp_audio_file.name)
        os.remove(tmp_video_file.name)
        for f in os.listdir(frames_dir):
            os.remove(os.path.join(frames_dir, f))
        os.rmdir(frames_dir)