Spaces:
Sleeping
Sleeping
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()
|