Function render

mbench/render.py:48–200

(motions, outdir='test_vis', device_id=0, name=None, render_video=False)

        return P

def render(motions, outdir='test_vis', device_id=0, name=None, render_video=False):
    frames, njoints, nfeats = motions.shape
    MINS = motions.min(axis=0).min(axis=0)
    MAXS = motions.max(axis=0).max(axis=0)

    height_offset = MINS[1]
    motions[:, :, 1] -= height_offset
    trajec = motions[:, 0, [0, 2]]

    j2s = joints2smpl(num_frames=frames, device_id=device_id, cuda=True)
    rot2xyz = Rotation2xyz(device=torch.device(f"cuda:{device_id}"))
    faces = rot2xyz.smpl_model.faces

    pt_output_path = os.path.join(outdir, f'{name}.pt')
    video_output_path = os.path.join(outdir, f'{name}.mp4')
    
    smplify_time = 0
    if not os.path.exists(pt_output_path): 
        print(f'Running SMPLify, it may take a few minutes.')
        smplify_start = time.time()
        
        pred_pose, motion_tensor, opt_dict = j2s.joint2smpl(motions)  # [nframes, njoints, 3]

        vertices = rot2xyz(motion_tensor.clone().detach(), mask=None,
                                        pose_rep='rot6d', translation=True, glob=True,
                                        jointstype='vertices',
                                        vertstrans=True)

        pred_pose_save = pred_pose.reshape(-1, 24, 3).cpu()
        vertices_save = vertices.squeeze(0).permute(2, 0, 1).cpu()
        data_to_save = {
            'pose': pred_pose_save,
            'joints': motions,
            'vertices': vertices_save
        }
        torch.save(data_to_save, pt_output_path)
        
        smplify_time = time.time() - smplify_start
        print(f'SMPLify completed in {smplify_time:.2f} seconds')
    else:
        print(f'Loading existing SMPLify results from {pt_output_path}')
        if render_video:
            vertices = torch.load(pt_output_path)['vertices']
            vertices = vertices.permute(1, 2, 0).unsqueeze(0)
    
    render_time = 0
    if render_video:
        print(f'Starting video rendering for {frames} frames...')
        render_start = time.time()
        
        frames = vertices.shape[3]
        MINS = torch.min(torch.min(vertices[0], axis=0)[0], axis=1)[0]
        MAXS = torch.max(torch.max(vertices[0], axis=0)[0], axis=1)[0]

        # Pre-compute static elements
        minx = MINS[0] - 0.5
        maxx = MAXS[0] + 0.5
        minz = MINS[2] - 0.5