OpenMotionLab/MotionGPT

1import numpy as np2import imageio3import os4import argparse5from tqdm import tqdm6from .renderer import get_renderer7 8 9def get_rotation(theta=np.pi / 3):10    import mGPT.utils.rotation_conversions as geometry11    import torch12    axis = torch.tensor([0, 1, 0], dtype=torch.float)13    axisangle = theta * axis14    matrix = geometry.axis_angle_to_matrix(axisangle)15    return matrix.numpy()16 17 18def render_video(meshes,19                 key,20                 action,21                 renderer,22                 savepath,23                 backgrounds,24                 cam_pose,25                 cams=(0.75, 0.75, 0, 0.10),26                 color=[0.11, 0.53, 0.8]):27    #  cams=(0.75, 0.75, 0, 0.10), color=[165.0/255,112/255,140/255]):28    # center the first frame29    if key not in ["real", "ntf", "side"]:30        w = int(key) / 6.031        # purpole to green32        # color = w*np.array([0.9,102/255,120/255]) + (1-w)*np.array([0.11, 0.9, 0.11])33        # color = (1-w)*np.array([165.0/255,112/255,140/255]) + w*np.array([0.11, 0.8, 0.11])34        color = (1 - w) * np.array([0.75, 0.13, 0.7]) + w * np.array(35            [0.12, 0.7, 0.14])36 37    meshes = meshes - meshes[0].mean(axis=0)38    imgs = []39    idx = 040    # for mesh in meshes:41    for mesh in tqdm(meshes, desc=f"Visualize {key}, action {action}"):42        # file_name = '3dpw_rot-90_glob_trimesh.ply' mesh_filename=file_name,43        # prepare background44        if len(backgrounds.shape) == 3:45            background = backgrounds46            cam = cams47        elif len(backgrounds.shape) == 4:48            background = backgrounds[idx]49            cam = cams[idx]50            idx += 151        # prepare cams52        img = renderer.render(background,53                              mesh,54                              cam,55                              color=color,56                              cam_pose=cam_pose)57        imgs.append(img)58        # show(img)59 60    imgs = np.array(imgs)61    # masks = ~(imgs/255. > 0.96).all(-1)62    # coords = np.argwhere(masks.sum(axis=0))63    # y1, x1 = coords.min(axis=0)64    # y2, x2 = coords.max(axis=0)65    # writer = imageio.get_writer(savepath, fps=30)66    # for cimg in imgs[:, y1:y2, x1:x2]:67    #     writer.append_data(cimg)68    # writer.close()69 70    # from mld.utils.uicap_utils import write_rgba_seqs71    # write_rgba_seqs(imgs, savepath)72 73    writer = imageio.get_writer(savepath, fps=30)74    for cimg in imgs:75        writer.append_data(cimg)76    writer.close()77 78 79def main():80    parser = argparse.ArgumentParser()81    parser.add_argument("filename")82    opt = parser.parse_args()83    filename = opt.filename84    savefolder = os.path.splitext(filename)[0]85    os.makedirs(savefolder, exist_ok=True)86 87    output = np.load(filename)88 89    if output.shape[0] == 3:90        visualization, generation, reconstruction = output91        output = {92            "visualization": visualization,93            "generation": generation,94            "reconstruction": reconstruction95        }96    else:97        # output = {f"generation_{key}": output[key] for key in range(2)} #  len(output))}98        # output = {f"generation_{key}": output[key] for key in range(len(output))}99        output = {100            f"generation_{key}": output[key]101            for key in range(len(output))102        }103 104    width = 1024105    height = 1024106 107    background = np.zeros((height, width, 3))108    renderer = get_renderer(width, height)109 110    # if duration mode, put back durations111    if output["generation_3"].shape[-1] == 100:112        output["generation_0"] = output["generation_0"][:, :, :, :40]113        output["generation_1"] = output["generation_1"][:, :, :, :60]114        output["generation_2"] = output["generation_2"][:, :, :, :80]115        output["generation_3"] = output["generation_3"][:, :, :, :100]116    elif output["generation_3"].shape[-1] == 160:117        print("160 mode")118        output["generation_0"] = output["generation_0"][:, :, :, :100]119        output["generation_1"] = output["generation_1"][:, :, :, :120]120        output["generation_2"] = output["generation_2"][:, :, :, :140]121        output["generation_3"] = output["generation_3"][:, :, :, :160]122 123    # if str(action) == str(1) and str(key) == "generation_4":124    for key in output:125        vidmeshes = output[key]126        for action in range(len(vidmeshes)):127            meshes = vidmeshes[action].transpose(2, 0, 1)128            path = os.path.join(savefolder,129                                "action{}_{}.mp4".format(action, key))130            render_video(meshes, key, action, renderer, path, background)131 132 133if __name__ == "__main__":134    main()135