File size: 4,988 Bytes
e7b9fb6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
136
# Modified from https://github.com/lab4d-org/lab4d

import os
import numpy as np
import cv2
import pyrender
import trimesh
from pyrender import (
    IntrinsicsCamera,
    Mesh,
    Node,
    Scene,
    OffscreenRenderer,
    MetallicRoughnessMaterial,
    RenderFlags
)

os.environ["PYOPENGL_PLATFORM"] = "egl"

def look_at(eye, center, up):
    """Create a look-at (view) matrix."""
    f = np.array(center, dtype=np.float32) - np.array(eye, dtype=np.float32)
    f /= np.linalg.norm(f)

    u = np.array(up, dtype=np.float32)
    u /= np.linalg.norm(u)

    s = np.cross(f, u)
    u = np.cross(s, f)

    m = np.identity(4, dtype=np.float32)
    m[0, :3] = s
    m[1, :3] = u
    m[2, :3] = -f
    m[:3, 3] = -np.matmul(m[:3, :3], np.array(eye, dtype=np.float32))

    return m

class PyRenderWrapper:
    def __init__(self, image_size=(1024, 1024)) -> None:
        # renderer
        self.image_size = image_size
        render_size = max(image_size)
        self.r = OffscreenRenderer(render_size, render_size)
        self.intrinsics = IntrinsicsCamera(
            render_size, render_size, render_size / 2, render_size / 2
        )
        # light
        self.light_pose = np.eye(4)
        self.set_light_topdown()
        self.direc_l = pyrender.DirectionalLight(color=np.ones(3), intensity=5.0)
        self.material = MetallicRoughnessMaterial(
            roughnessFactor=0.75, metallicFactor=0.75, alphaMode="BLEND"
        )
        self.init_camera()

    def init_camera(self):
        self.flip_pose = np.eye(4)
        self.set_camera(np.eye(4))

    def set_camera(self, scene_to_cam):
        # object to camera transforms
        self.scene_to_cam = self.flip_pose @ scene_to_cam
        
    def set_light_topdown(self, gl=False):
        # top down light, slightly closer to the camera
        if gl:
            rot = cv2.Rodrigues(np.asarray([-np.pi / 2, 0, 0]))[0]
        else:
            rot = cv2.Rodrigues(np.asarray([np.pi / 2, 0, 0]))[0]
        self.light_pose[:3, :3] = rot

    def align_light_to_camera(self):
        self.light_pose = np.linalg.inv(self.scene_to_cam)

    def set_intrinsics(self, intrinsics):
        """

        Args:

            intrinsics: (4,) fx,fy,px,py

        """
        self.intrinsics = IntrinsicsCamera(
            intrinsics[0], intrinsics[1], intrinsics[2], intrinsics[3]
        )

    def get_cam_to_scene(self):
        cam_to_scene = np.eye(4)
        cam_to_scene[:3, :3] = self.scene_to_cam[:3, :3].T
        cam_to_scene[:3, 3] = -self.scene_to_cam[:3, :3].T @ self.scene_to_cam[:3, 3]
        return cam_to_scene
    
    def set_camera_view(self, angle, bbox_center, distance=2.0):
        # Calculate camera position based on angle and distance from bounding box center
        camera_position = bbox_center + distance * np.array([np.sin(angle), 0, np.cos(angle)], dtype=np.float32)
        look_at_matrix = look_at(camera_position, bbox_center, [0, 1, 0])
        self.scene_to_cam = look_at_matrix @ self.flip_pose
    
    def render(self, input_dict):
        # Create separate scenes for transparent objects (mesh) and solid objects (joints and bones)
        scene_transparent = Scene(ambient_light=np.array([1.0, 1.0, 1.0, 1.0]) * 0.1)
        scene_solid = Scene(ambient_light=np.array([1.0, 1.0, 1.0, 1.0]) * 0.1)

        mesh_pyrender = Mesh.from_trimesh(input_dict["shape"], smooth=False)
        mesh_pyrender.primitives[0].material = self.material
        scene_transparent.add(mesh_pyrender, pose=np.eye(4), name="shape")

        if "joint_meshes" in input_dict:
            joints_pyrender = Mesh.from_trimesh(input_dict["joint_meshes"], smooth=False)
            joints_pyrender.primitives[0].material = self.material
            scene_solid.add(joints_pyrender, pose=np.eye(4), name="joints")

        if "bone_meshes" in input_dict:
            bones_pyrender = Mesh.from_trimesh(input_dict["bone_meshes"], smooth=False)
            bones_pyrender.primitives[0].material = self.material
            scene_solid.add(bones_pyrender, pose=np.eye(4), name="bones")

        # Camera for both scenes
        scene_transparent.add(self.intrinsics, pose=self.get_cam_to_scene())
        scene_solid.add(self.intrinsics, pose=self.get_cam_to_scene())

        # Light for both scenes
        scene_transparent.add(self.direc_l, pose=self.light_pose)
        scene_solid.add(self.direc_l, pose=self.light_pose)

        # Render transparent scene first
        color_transparent, depth_transparent = self.r.render(scene_transparent)

        # Render solid scene on top
        color_solid, depth_solid = self.r.render(scene_solid)

        # Combine the two scenes
        color_combined = np.where(depth_solid[..., np.newaxis] == 0, color_transparent, color_solid)

        return color_combined, depth_solid
    def delete(self):
        self.r.delete()