File size: 2,771 Bytes
2d5fdd1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7a7e0f4
 
2d5fdd1
 
 
 
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
from smplx import SMPL
import torch
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import os
import glob
root_dir = os.path.dirname(os.path.realpath(__file__))
from analysis.visualization.convert_mat_to_euler import rot_mats_to_eulers
from .utils import generate_video_from_images, join_video_and_audio

def delete_images():
    files = glob.glob(root_dir+'/img/*')
    for f in files:
        os.remove(f)

def generate_video_from_mats(pred_mats_file, output_folder, audio_file, trim_audio=0, fps=60, plot_mats=False):
    filename = os.path.basename(pred_mats_file)
    seq_id = filename.split(".")[0]
    predicted_mats = np.load(pred_mats_file)
    if plot_mats:
        plt.matshow(predicted_mats[:1000,0,:-3])
        plt.savefig(root_dir+"/last.predicted_mats.png")

    smpl_thing = rot_mats_to_eulers(predicted_mats)
    smpl_poses,smpl_scaling,smpl_trans = smpl_thing['smpl_poses'], smpl_thing['smpl_scaling'], smpl_thing['smpl_trans']
    #this loads the SMPL_MALE pickle
    smpl = SMPL(model_path="./aistplusplus_api", gender='MALE', batch_size=1) #what if we try gender='FEMALE'?
    output = smpl.forward(
        global_orient=torch.from_numpy(smpl_poses[:, 0:1]).float(),
        body_pose=torch.from_numpy(smpl_poses[:, 1:]).float(),
        transl=torch.from_numpy(smpl_trans).float(),
        scaling=torch.from_numpy(smpl_scaling.reshape(1, 1)).float(),
        )
    keypoints3d = output.joints.detach().numpy()
    keypoints3d = keypoints3d[:,:24] # the body joints (ignoring the extra head, feet and hand bones added onto it here https://github.com/vchoutas/smplx/blob/7547ee6656b942a68a97604d0cf7b6b834fad9eb/smplx/vertex_joint_selector.py)
    # that file takes the position of the vertices corresponding to certain joints
    # print(keypoints3d)

    # Plot as images
    delete_images()
    fig = plt.figure()
    plt.ion()
    plt.show()
    ax = Axes3D(fig)
    ax.scatter(keypoints3d[0,:,2], keypoints3d[0,:,0], keypoints3d[0,:,1])
    plt.xlim([-100,100])
    plt.ylim([-100,100])
    ax.set_zlim([75,275])
    ax.view_init(0, 0)
    plt.draw()
    plt.pause(0.001)

    for i in range(1,len(keypoints3d)):
        print(i)
        ax.clear()
        ax.scatter(keypoints3d[i,:,2], keypoints3d[i,:,0], keypoints3d[i,:,1])
        plt.xlim([-100,100])
        plt.ylim([-100,100])
        ax.set_zlim([75,275])
        ax.view_init(0, 0)
        plt.draw()
        plt.pause(0.001)
        plt.savefig(root_dir+"/img/img_"+str(i)+".png")

    #video_file = output_folder+seq_id+".mp4"
    video_file = f'{output_folder}/example.mp4'
    generate_video_from_images(root_dir+"/img", video_file, fps)

    if audio_file is not None:
        join_video_and_audio(video_file, audio_file, trim_audio)