File size: 4,602 Bytes
4409449
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
import numpy as np
import imageio
import os
import argparse
from tqdm import tqdm
from .renderer import get_renderer


def get_rotation(theta=np.pi / 3):
    import mGPT.utils.rotation_conversions as geometry
    import torch
    axis = torch.tensor([0, 1, 0], dtype=torch.float)
    axisangle = theta * axis
    matrix = geometry.axis_angle_to_matrix(axisangle)
    return matrix.numpy()


def render_video(meshes,
                 key,
                 action,
                 renderer,
                 savepath,
                 backgrounds,
                 cam_pose,
                 cams=(0.75, 0.75, 0, 0.10),
                 color=[0.11, 0.53, 0.8]):
    #  cams=(0.75, 0.75, 0, 0.10), color=[165.0/255,112/255,140/255]):
    # center the first frame
    if key not in ["real", "ntf", "side"]:
        w = int(key) / 6.0
        # purpole to green
        # color = w*np.array([0.9,102/255,120/255]) + (1-w)*np.array([0.11, 0.9, 0.11])
        # color = (1-w)*np.array([165.0/255,112/255,140/255]) + w*np.array([0.11, 0.8, 0.11])
        color = (1 - w) * np.array([0.75, 0.13, 0.7]) + w * np.array(
            [0.12, 0.7, 0.14])

    meshes = meshes - meshes[0].mean(axis=0)
    imgs = []
    idx = 0
    # for mesh in meshes:
    for mesh in tqdm(meshes, desc=f"Visualize {key}, action {action}"):
        # file_name = '3dpw_rot-90_glob_trimesh.ply' mesh_filename=file_name,
        # prepare background
        if len(backgrounds.shape) == 3:
            background = backgrounds
            cam = cams
        elif len(backgrounds.shape) == 4:
            background = backgrounds[idx]
            cam = cams[idx]
            idx += 1
        # prepare cams
        img = renderer.render(background,
                              mesh,
                              cam,
                              color=color,
                              cam_pose=cam_pose)
        imgs.append(img)
        # show(img)

    imgs = np.array(imgs)
    # masks = ~(imgs/255. > 0.96).all(-1)
    # coords = np.argwhere(masks.sum(axis=0))
    # y1, x1 = coords.min(axis=0)
    # y2, x2 = coords.max(axis=0)
    # writer = imageio.get_writer(savepath, fps=30)
    # for cimg in imgs[:, y1:y2, x1:x2]:
    #     writer.append_data(cimg)
    # writer.close()

    # from mld.utils.uicap_utils import write_rgba_seqs
    # write_rgba_seqs(imgs, savepath)

    writer = imageio.get_writer(savepath, fps=30)
    for cimg in imgs:
        writer.append_data(cimg)
    writer.close()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("filename")
    opt = parser.parse_args()
    filename = opt.filename
    savefolder = os.path.splitext(filename)[0]
    os.makedirs(savefolder, exist_ok=True)

    output = np.load(filename)

    if output.shape[0] == 3:
        visualization, generation, reconstruction = output
        output = {
            "visualization": visualization,
            "generation": generation,
            "reconstruction": reconstruction
        }
    else:
        # output = {f"generation_{key}": output[key] for key in range(2)} #  len(output))}
        # output = {f"generation_{key}": output[key] for key in range(len(output))}
        output = {
            f"generation_{key}": output[key]
            for key in range(len(output))
        }

    width = 1024
    height = 1024

    background = np.zeros((height, width, 3))
    renderer = get_renderer(width, height)

    # if duration mode, put back durations
    if output["generation_3"].shape[-1] == 100:
        output["generation_0"] = output["generation_0"][:, :, :, :40]
        output["generation_1"] = output["generation_1"][:, :, :, :60]
        output["generation_2"] = output["generation_2"][:, :, :, :80]
        output["generation_3"] = output["generation_3"][:, :, :, :100]
    elif output["generation_3"].shape[-1] == 160:
        print("160 mode")
        output["generation_0"] = output["generation_0"][:, :, :, :100]
        output["generation_1"] = output["generation_1"][:, :, :, :120]
        output["generation_2"] = output["generation_2"][:, :, :, :140]
        output["generation_3"] = output["generation_3"][:, :, :, :160]

    # if str(action) == str(1) and str(key) == "generation_4":
    for key in output:
        vidmeshes = output[key]
        for action in range(len(vidmeshes)):
            meshes = vidmeshes[action].transpose(2, 0, 1)
            path = os.path.join(savefolder,
                                "action{}_{}.mp4".format(action, key))
            render_video(meshes, key, action, renderer, path, background)


if __name__ == "__main__":
    main()