Spaces:
Running on Zero
Running on Zero
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: Apache-2.0 | |
| """Minimal pyrender-based mesh-on-image compositor. | |
| Drop-in for the subset of sam_3d_body.visualization.renderer.Renderer we use: | |
| Renderer(focal_length, faces).__call__(vertices, cam_t, image, camera_center=None) | |
| """ | |
| import os | |
| os.environ.setdefault("PYOPENGL_PLATFORM", "osmesa") # must precede first pyrender import | |
| import numpy as np | |
| import pyrender | |
| import trimesh | |
| def _create_raymond_lights(): | |
| thetas = np.pi * np.array([1.0 / 6.0, 1.0 / 6.0, 1.0 / 6.0]) | |
| phis = np.pi * np.array([0.0, 2.0 / 3.0, 4.0 / 3.0]) | |
| nodes = [] | |
| for phi, theta in zip(phis, thetas): | |
| xp = np.sin(theta) * np.cos(phi) | |
| yp = np.sin(theta) * np.sin(phi) | |
| zp = np.cos(theta) | |
| z = np.array([xp, yp, zp]) | |
| z /= np.linalg.norm(z) | |
| x = np.array([-z[1], z[0], 0.0]) | |
| if np.linalg.norm(x) == 0: | |
| x = np.array([1.0, 0.0, 0.0]) | |
| x /= np.linalg.norm(x) | |
| y = np.cross(z, x) | |
| matrix = np.eye(4) | |
| matrix[:3, :3] = np.c_[x, y, z] | |
| nodes.append( | |
| pyrender.Node( | |
| light=pyrender.DirectionalLight(color=np.ones(3), intensity=1.0), | |
| matrix=matrix, | |
| ) | |
| ) | |
| return nodes | |
| class Renderer: | |
| def __init__(self, focal_length, faces=None): | |
| self.focal_length = focal_length | |
| self.faces = faces | |
| def __call__( | |
| self, | |
| vertices: np.ndarray, # (V, 3) | |
| cam_t: np.ndarray, # (3,) | |
| image: np.ndarray, # (H, W, 3) uint8 | |
| mesh_base_color=(1.0, 1.0, 0.9), | |
| camera_center=None, | |
| ) -> np.ndarray: # (H, W, 3) uint8 | |
| image_f = image.astype(np.float32) / 255.0 | |
| h, w = image_f.shape[:2] | |
| if camera_center is None: | |
| camera_center = [w / 2.0, h / 2.0] | |
| camera_translation = cam_t.copy() | |
| camera_translation[0] *= -1.0 | |
| material = pyrender.MetallicRoughnessMaterial( | |
| metallicFactor=0.0, | |
| alphaMode="OPAQUE", | |
| baseColorFactor=(mesh_base_color[2], mesh_base_color[1], mesh_base_color[0], 1.0), | |
| ) | |
| mesh = trimesh.Trimesh(vertices.copy(), self.faces.copy()) | |
| rot = trimesh.transformations.rotation_matrix(np.radians(180), [1, 0, 0]) | |
| mesh.apply_transform(rot) | |
| mesh = pyrender.Mesh.from_trimesh(mesh, material=material) | |
| scene = pyrender.Scene(bg_color=[0, 0, 0, 0], ambient_light=(0.3, 0.3, 0.3)) | |
| scene.add(mesh, "mesh") | |
| camera_pose = np.eye(4) | |
| camera_pose[:3, 3] = camera_translation | |
| camera = pyrender.IntrinsicsCamera( | |
| fx=self.focal_length, | |
| fy=self.focal_length, | |
| cx=camera_center[0], | |
| cy=camera_center[1], | |
| zfar=1e12, | |
| ) | |
| scene.add(camera, pose=camera_pose) | |
| for node in _create_raymond_lights(): | |
| scene.add_node(node) | |
| r = pyrender.OffscreenRenderer(viewport_width=w, viewport_height=h) | |
| color, depth = r.render(scene, flags=pyrender.RenderFlags.RGBA) | |
| r.delete() | |
| color_f = color.astype(np.float32) / 255.0 | |
| if color_f.shape[2] == 4: | |
| alpha = color_f[:, :, 3:4] | |
| out = color_f[:, :, :3] * alpha + (1.0 - alpha) * image_f | |
| else: | |
| # RGBA not supported by this backend; composite via depth mask | |
| mask = (depth > 0)[..., np.newaxis] | |
| out = np.where(mask, color_f, image_f) | |
| return (out * 255).clip(0, 255).astype(np.uint8) | |