vumichien/Generate_human_motion

1from models.rotation2xyz import Rotation2xyz2import numpy as np3from trimesh import Trimesh4import os5os.environ['PYOPENGL_PLATFORM'] = "osmesa"6 7import torch8from visualize.simplify_loc2rot import joints2smpl9import pyrender10import matplotlib.pyplot as plt11 12import io13import imageio14from shapely import geometry15import trimesh16from pyrender.constants import RenderFlags17import math18# import ffmpeg19from PIL import Image20 21class WeakPerspectiveCamera(pyrender.Camera):22    def __init__(self,23                 scale,24                 translation,25                 znear=pyrender.camera.DEFAULT_Z_NEAR,26                 zfar=None,27                 name=None):28        super(WeakPerspectiveCamera, self).__init__(29            znear=znear,30            zfar=zfar,31            name=name,32        )33        self.scale = scale34        self.translation = translation35 36    def get_projection_matrix(self, width=None, height=None):37        P = np.eye(4)38        P[0, 0] = self.scale[0]39        P[1, 1] = self.scale[1]40        P[0, 3] = self.translation[0] * self.scale[0]41        P[1, 3] = -self.translation[1] * self.scale[1]42        P[2, 2] = -143        return P44 45def render(motions, outdir='test_vis', device_id=0, name=None, pred=True):46    frames, njoints, nfeats = motions.shape47    MINS = motions.min(axis=0).min(axis=0)48    MAXS = motions.max(axis=0).max(axis=0)49 50    height_offset = MINS[1]51    motions[:, :, 1] -= height_offset52    trajec = motions[:, 0, [0, 2]]53 54    j2s = joints2smpl(num_frames=frames, device_id=0, cuda=True)55    rot2xyz = Rotation2xyz(device=torch.device("cuda:0"))56    faces = rot2xyz.smpl_model.faces57 58    if (not os.path.exists(outdir + name+'_pred.pt') and pred) or (not os.path.exists(outdir + name+'_gt.pt') and not pred): 59        print(f'Running SMPLify, it may take a few minutes.')60        motion_tensor, opt_dict = j2s.joint2smpl(motions)  # [nframes, njoints, 3]61 62        vertices = rot2xyz(torch.tensor(motion_tensor).clone(), mask=None,63                                        pose_rep='rot6d', translation=True, glob=True,64                                        jointstype='vertices',65                                        vertstrans=True)66 67        if pred:68            torch.save(vertices, outdir + name+'_pred.pt')69        else:70            torch.save(vertices, outdir + name+'_gt.pt')71    else:72        if pred:73            vertices = torch.load(outdir + name+'_pred.pt')74        else:75            vertices = torch.load(outdir + name+'_gt.pt')76    frames = vertices.shape[3] # shape: 1, nb_frames, 3, nb_joints77    print (vertices.shape)78    MINS = torch.min(torch.min(vertices[0], axis=0)[0], axis=1)[0]79    MAXS = torch.max(torch.max(vertices[0], axis=0)[0], axis=1)[0]80    # vertices[:,:,1,:] -= MINS[1] + 1e-581 82 83    out_list = []84    85    minx = MINS[0] - 0.586    maxx = MAXS[0] + 0.587    minz = MINS[2] - 0.5 88    maxz = MAXS[2] + 0.589    polygon = geometry.Polygon([[minx, minz], [minx, maxz], [maxx, maxz], [maxx, minz]])90    polygon_mesh = trimesh.creation.extrude_polygon(polygon, 1e-5)91 92    vid = []93    for i in range(frames):94        if i % 10 == 0:95            print(i)96 97        mesh = Trimesh(vertices=vertices[0, :, :, i].squeeze().tolist(), faces=faces)98 99        base_color = (0.11, 0.53, 0.8, 0.5)100        ## OPAQUE rendering without alpha101        ## BLEND rendering consider alpha 102        material = pyrender.MetallicRoughnessMaterial(103            metallicFactor=0.7,104            alphaMode='OPAQUE',105            baseColorFactor=base_color106        )107 108 109        mesh = pyrender.Mesh.from_trimesh(mesh, material=material)110 111        polygon_mesh.visual.face_colors = [0, 0, 0, 0.21]112        polygon_render = pyrender.Mesh.from_trimesh(polygon_mesh, smooth=False)113 114        bg_color = [1, 1, 1, 0.8]115        scene = pyrender.Scene(bg_color=bg_color, ambient_light=(0.4, 0.4, 0.4))116        117        sx, sy, tx, ty = [0.75, 0.75, 0, 0.10]118 119        camera = pyrender.PerspectiveCamera(yfov=(np.pi / 3.0))120 121        light = pyrender.DirectionalLight(color=[1,1,1], intensity=300)122 123        scene.add(mesh)124 125        c = np.pi / 2126 127        scene.add(polygon_render, pose=np.array([[ 1, 0, 0, 0],128 129        [ 0, np.cos(c), -np.sin(c), MINS[1].cpu().numpy()],130 131        [ 0, np.sin(c), np.cos(c), 0],132 133        [ 0, 0, 0, 1]]))134 135        light_pose = np.eye(4)136        light_pose[:3, 3] = [0, -1, 1]137        scene.add(light, pose=light_pose.copy())138 139        light_pose[:3, 3] = [0, 1, 1]140        scene.add(light, pose=light_pose.copy())141 142        light_pose[:3, 3] = [1, 1, 2]143        scene.add(light, pose=light_pose.copy())144 145 146        c = -np.pi / 6147 148        scene.add(camera, pose=[[ 1, 0, 0, (minx+maxx).cpu().numpy()/2],149 150                                [ 0, np.cos(c), -np.sin(c), 1.5],151 152                                [ 0, np.sin(c), np.cos(c), max(4, minz.cpu().numpy()+(1.5-MINS[1].cpu().numpy())*2, (maxx-minx).cpu().numpy())],153 154                                [ 0, 0, 0, 1]155                                ])156        157        # render scene158        r = pyrender.OffscreenRenderer(960, 960)159 160        color, _ = r.render(scene, flags=RenderFlags.RGBA)161        # Image.fromarray(color).save(outdir+'/'+name+'_'+str(i)+'.png')162 163        vid.append(color)164 165        r.delete()166 167    out = np.stack(vid, axis=0)168    if pred:169        imageio.mimsave(outdir + name+'_pred.gif', out, fps=20)170    else:171        imageio.mimsave(outdir + name+'_gt.gif', out, fps=20)172    173 174 175 176 177if __name__ == "__main__":178    import argparse179    parser = argparse.ArgumentParser()180    parser.add_argument("--filedir", type=str, default=None, help='motion npy file dir')181    parser.add_argument('--motion-list', default=None, nargs="+", type=str, help="motion name list")182    args = parser.parse_args()183 184    filename_list = args.motion_list185    filedir = args.filedir186    187    for filename in filename_list:188        motions = np.load(filedir + filename+'_pred.npy')189        print('pred', motions.shape, filename)190        render(motions[0], outdir=filedir, device_id=0, name=filename, pred=True)191 192        motions = np.load(filedir + filename+'_gt.npy')193        print('gt', motions.shape, filename)194        render(motions[0], outdir=filedir, device_id=0, name=filename, pred=False)195