OpenMotionLab/MotionGPT

1"""2This script is borrowed from https://github.com/mkocabas/VIBE3 Adhere to their licence to use this script4 It has been modified5"""6 7import os8import math9import trimesh10 11import pyrender12import numpy as np13from pyrender.constants import RenderFlags14 15 16# os.environ['DISPLAY'] = ':0.0'17# os.environ['PYOPENGL_PLATFORM'] = 'egl'18# os.environ['PYOPENGL_PLATFORM'] = 'osmesa'19SMPL_MODEL_DIR = "data/smpl_data/"20 21 22def get_smpl_faces():23    return np.load(os.path.join(SMPL_MODEL_DIR, "smplfaces.npy"))24 25 26class WeakPerspectiveCamera(pyrender.Camera):27    def __init__(self,28                 scale,29                 translation,30                 znear=pyrender.camera.DEFAULT_Z_NEAR,31                 zfar=None,32                 name=None):33        super(WeakPerspectiveCamera, self).__init__(34            znear=znear,35            zfar=zfar,36            name=name,37        )38        self.scale = scale39        self.translation = translation40 41    def get_projection_matrix(self, width=None, height=None):42        P = np.eye(4)43        P[0, 0] = self.scale[0]44        P[1, 1] = self.scale[1]45        P[0, 3] = self.translation[0] * self.scale[0]46        P[1, 3] = -self.translation[1] * self.scale[1]47        P[2, 2] = -148        return P49 50 51class Renderer:52    def __init__(self, background=None, resolution=(224, 224), bg_color=[0, 0, 0, 0.5], orig_img=False, wireframe=False, cam_pose=np.eye(4)):53        width, height = resolution54        self.background = np.zeros((height, width, 3))55        self.resolution = resolution56 57        self.faces = get_smpl_faces()58        self.orig_img = orig_img59        self.wireframe = wireframe60        self.renderer = pyrender.OffscreenRenderer(61            viewport_width=self.resolution[0],62            viewport_height=self.resolution[1],63            point_size=0.564        )65 66        # set the scene67        self.scene = pyrender.Scene(bg_color=bg_color, ambient_light=(0.4, 0.4, 0.4))68 69        light = pyrender.PointLight(color=[1.0, 1.0, 1.0], intensity=4)70    71 72        light_pose = np.eye(4)73        light_pose[:3, 3] = [0, -1, 1]74        self.scene.add(light, pose=np.dot(cam_pose,light_pose).copy())75 76        light_pose[:3, 3] = [0, 1, 1]77        self.scene.add(light, pose=np.dot(cam_pose,light_pose).copy())78 79        light_pose[:3, 3] = [1, 1, 2]80        self.scene.add(light, pose=np.dot(cam_pose,light_pose).copy())81 82        """ok83        light_pose = np.eye(4)84        light_pose[:3, 3] = [0, -1, 1]85        self.scene.add(light, pose=light_pose)86 87        light_pose[:3, 3] = [0, 1, 1]88        self.scene.add(light, pose=light_pose)89 90        light_pose[:3, 3] = [1, 1, 2]91        self.scene.add(light, pose=light_pose)92        """93 94        # light_pose[:3, 3] = [0, -2, 2]95        # [droite, hauteur, profondeur camera]96        """97        light_pose = np.eye(4)98        light_pose[:3, 3] = [0, -1, 1]99        self.scene.add(light, pose=light_pose)100 101        light_pose[:3, 3] = [0, 1, 1]102        self.scene.add(light, pose=light_pose)103 104        light_pose[:3, 3] = [1, 1, 2]105        self.scene.add(light, pose=light_pose)106        """107 108    def render(self, img, verts, cam, angle=None, axis=None, mesh_filename=None, color=[1.0, 1.0, 0.9],109               cam_pose=np.eye(4)):110        mesh = trimesh.Trimesh(vertices=verts, faces=self.faces, process=False)111        Rx = trimesh.transformations.rotation_matrix(math.radians(180), [1, 0, 0])112        # Rx = trimesh.transformations.rotation_matrix(math.radians(-90), [1, 0, 0])113        mesh.apply_transform(Rx)114 115        if mesh_filename is not None:116            mesh.export(mesh_filename)117 118        if angle and axis:119            R = trimesh.transformations.rotation_matrix(math.radians(angle), axis)120            mesh.apply_transform(R)121 122        sx, sy, tx, ty = cam123 124        camera = WeakPerspectiveCamera(125            scale=[sx, sy],126            translation=[tx, ty],127            zfar=100000.128        )129 130        material = pyrender.MetallicRoughnessMaterial(131            metallicFactor=0.0, # 0.0 for no specular lighting132            # metallicFactor=0.7, # 0.0 for no specular lighting133            alphaMode='OPAQUE',134            baseColorFactor=(color[0], color[1], color[2], 1.0)135        )136 137        mesh = pyrender.Mesh.from_trimesh(mesh, material=material)138 139        mesh_node = self.scene.add(mesh, 'mesh')140 141        cam_node = self.scene.add(camera, pose=cam_pose)142 143        if self.wireframe:144            render_flags = RenderFlags.RGBA | RenderFlags.ALL_WIREFRAME145        else:146            render_flags = RenderFlags.RGBA147 148        rgb, _ = self.renderer.render(self.scene, flags=render_flags)149        if rgb.shape[-1]==3:150            # Debug151            # 0 not distinguish alpha152            valid_mask = (rgb[:, :, -1] > 0)[:, :, np.newaxis]153            output_img = rgb * valid_mask + (1 - valid_mask) * img154        elif rgb.shape[-1]==4:155            # valid_mask = (rgb[:, :, -1] > 128)[:, :, np.newaxis]156            # output_img = rgb[:, :, :-1] * valid_mask + (1 - valid_mask) * img157 158            # # output alpha159            valid_mask = (rgb[:, :, -1] > 128)[:, :]160            output_img = np.copy(rgb)161            output_img[:, :, -1] *= valid_mask           162            # output_img = img163        else:164            raise ValueError(f"rgb shape {rgb.shape[-1]} is not correct!")165        image = output_img.astype(np.uint8)166 167        self.scene.remove_node(mesh_node)168        self.scene.remove_node(cam_node)169 170        return image171 172 173def get_renderer(width, height, cam_pose):174    renderer = Renderer(resolution=(width, height),175                        bg_color=[1, 1, 1, 0.5],176                        orig_img=False,177                        wireframe=False,178                        cam_pose=cam_pose)179    return renderer180