File size: 4,233 Bytes
7b95dc2 | 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 137 138 139 140 141 142 143 144 145 146 147 | import copy
import numpy as np
def center_normalize(pcl):
# Center the point cloud
centroid = np.mean(pcl, axis=0)
pcl_centered = pcl - centroid
# Scale
max_dist = np.max(pcl_centered)
pcl_normalized = pcl_centered / (2 * max_dist)
# Move down till one the lowest point is on z=-0.5
pcl_normalized[:, 2] += -0.5 - np.min(pcl_normalized[:, 2])
return pcl_normalized.astype(np.float32)
try:
import pyrender
import trimesh
except ImportError: # Mitsuba-only 脚本只需 center_normalize，可不装 pyrender
pyrender = None
trimesh = None
def pretty_color(x, y, z):
x += 0.5
y += +0.5
z = z + 0.5 - 0.0125
vec = np.array([x, y, z])
vec = np.clip(vec, 0.001, 1.0)
norm = np.sqrt(np.sum(vec ** 2))
vec /= norm
return [vec[0], vec[1], vec[2]]
# code from https://github.com/nv-nguyen/cnos/blob/main/src/poses/pyrender.py
class Render(object):
def __init__(self, light_intensity):
if pyrender is None:
raise ImportError("Render 类需要安装 pyrender、trimesh：pip install pyrender trimesh")
# camera pose is fixed as np.eye(4)
cam_pose = np.eye(4)
# convert openCV camera
cam_pose[1, 1] = -1
cam_pose[2, 2] = -1
self.cam_pose = cam_pose
# create scene config
self.ambient_light = np.array(
[0.3, 0.3, 0.3, 0.3]) # np.array([0.2, 0.2, 0.2]) # np.array([1.0, 1.0, 1.0, 1.0])
self.light = pyrender.SpotLight(
color=np.ones(3),
intensity=light_intensity,
innerConeAngle=np.pi / 2.0,
outerConeAngle=np.pi / 2.0,
)
self.render_engine = {}
def render(
self,
mesh,
obj_poses,
img_size,
intrinsic,
dyn_obj_scale=False,
factor=1.5,
):
scene = pyrender.Scene(
bg_color=np.array([0., 0., 0., 1.0]), ambient_light=self.ambient_light
)
scene.add(self.light, pose=self.cam_pose)
# create camera and render engine
fx, fy, cx, cy = intrinsic[0], intrinsic[4], intrinsic[2], intrinsic[5]
camera = pyrender.IntrinsicsCamera(
fx=fx, fy=fy, cx=cx, cy=cy, znear=0.05, zfar=100000
)
scene.add(camera, pose=self.cam_pose)
if not dyn_obj_scale:
cad_node = scene.add(mesh, pose=np.eye(4), name="cad")
if (img_size[1], img_size[0]) not in self.render_engine:
self.render_engine[(img_size[1], img_size[0])] = pyrender.OffscreenRenderer(img_size[1], img_size[0])
rgbs, depths, masks = [], [], []
for idx_frame in range(obj_poses.shape[0]):
if dyn_obj_scale:
trans = obj_poses[idx_frame][:3, 3]
dist = ((trans ** 2).sum() ** 0.5) * factor
obj_poses[idx_frame][:3, 3] /= dist
scaled_mesh = copy.deepcopy(mesh).apply_scale(1 / dist)
pymesh = pyrender.Mesh.from_trimesh(as_mesh(scaled_mesh), smooth=False)
cad_node = scene.add(pymesh, pose=np.eye(4), name="cad")
scene.set_pose(cad_node, obj_poses[idx_frame])
rgb, depth = self.render_engine[(img_size[1], img_size[0])].render(scene,
pyrender.constants.RenderFlags.RGBA)
# img = Image.fromarray(np.uint8(rgb))
# img.save(output_dir + f"_{idx_frame:06d}.png")
if dyn_obj_scale:
depth *= dist
scene.remove_node(cad_node)
mask = depth > 0
masks.append(mask)
rgbs.append(rgb[..., :3])
depths.append(depth)
return rgbs, depths, masks
def as_mesh(scene_or_mesh):
if trimesh is None:
raise ImportError("as_mesh 需要 trimesh")
if isinstance(scene_or_mesh, trimesh.Scene):
result = trimesh.util.concatenate(
[
trimesh.Trimesh(vertices=m.vertices, faces=m.faces)
for m in scene_or_mesh.geometry.values()
]
)
else:
result = scene_or_mesh
return result
|