def add_object(self, augment_param=None):
for obj in self.obj_list[:]:
target_model_dir = os.path.join(
ENV_MODEL_DIR, obj.name) if obj.isEnv else os.path.join(
OBJ_MODEL_DIR, obj.name)
texture = cv2.cvtColor(
cv2.imread(target_model_dir + "/" + obj.name + ".png"),
cv2.COLOR_BGR2RGB)
texture = pyrender.Texture(source=texture, source_channels='RGB')
if augment_param:
texture = augment_param.augment_image(texture)
material = pyrender.MetallicRoughnessMaterial(
baseColorTexture=texture,
metallicFactor=0.2,
roughnessFactor=0.8)
obj_render_mesh = pyrender.Mesh.from_trimesh(obj.mesh,
material=material)
if self.in_workspace(obj.pose[:3, 3]) or obj.isEnv:
self.scene.add(obj_render_mesh,
pose=obj.pose,
name=str(obj.key))
else:
self.obj_list.remove(obj)
def render_skeleton(self, joints, camera_translation, image):
material = pyrender.MetallicRoughnessMaterial(
metallicFactor=0.2,
alphaMode='OPAQUE',
baseColorFactor=(0.8, 0.3, 0.3, 1.0))
camera_translation[0] *= -1.
scene = pyrender.Scene()
# print(joints.shape)
sm = trimesh.creation.uv_sphere(radius=0.01)
sm.visual.vertex_colors = [0.9, 0.1, 0.1, 1.0]
tfs = np.tile(np.eye(4), (len(joints), 1, 1))
tfs[:, :3, 3] = joints
tfs[:, 1, 1] = -1
tfs[:, 2, 2] = -1
rot = trimesh.transformations.rotation_matrix(np.radians(180),
[1, 0, 0])
# print(rot)
joints_pcl = pyrender.Mesh.from_trimesh(sm,
poses=tfs,
material=material)
scene.add(joints_pcl, 'joints')
# p = trimesh.load_path([[joints[29,:], joints[30,:]]])
bones = pyrender.Mesh.from_points([[joints[29, :], joints[30, :]]])
scene.add(bones, 'bones')
camera_pose = np.eye(4)
camera_pose[:3, 3] = camera_translation
camera = pyrender.IntrinsicsCamera(fx=self.focal_length,
fy=self.focal_length,
cx=self.camera_center[0],
cy=self.camera_center[1])
scene.add(camera, pose=camera_pose)
light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=1)
light_pose = np.eye(4)
light_pose[:3, 3] = np.array([0, -1, 1])
scene.add(light, pose=light_pose)
light_pose[:3, 3] = np.array([0, 1, 1])
scene.add(light, pose=light_pose)
light_pose[:3, 3] = np.array([1, 1, 2])
scene.add(light, pose=light_pose)
# pyrender.Viewer(scene)
color, rend_depth = self.renderer.render(
scene, flags=pyrender.RenderFlags.RGBA)
color = color.astype(np.float32) / 255.0
valid_mask = (rend_depth > 0)[:, :, None]
output_img = (color[:, :, :3] * valid_mask + (1 - valid_mask) * image)
return output_img
def render(self,
img,
verts,
cam,
angle=None,
axis=None,
mesh_filename=None,
color=[1.0, 1.0, 0.9]):
mesh = trimesh.Trimesh(vertices=verts, faces=self.faces)
Rx = trimesh.transformations.rotation_matrix(math.radians(180),
[1, 0, 0])
mesh.apply_transform(Rx)
if mesh_filename is not None:
mesh.export(mesh_filename)
if angle and axis:
R = trimesh.transformations.rotation_matrix(
math.radians(angle), axis)
mesh.apply_transform(R)
if cam.shape[-1] == 4:
sx, sy, tx, ty = cam
elif cam.shape[-1] == 3:
s, tx, ty = cam
sx = sy = s
camera = WeakPerspectiveCamera(scale=[sx, sy],
translation=[tx, ty],
zfar=1000.)
material = pyrender.MetallicRoughnessMaterial(
metallicFactor=0.0,
alphaMode='OPAQUE',
baseColorFactor=(color[0], color[1], color[2], 1.0))
mesh = pyrender.Mesh.from_trimesh(mesh, material=material)
mesh_node = self.scene.add(mesh, 'mesh')
camera_pose = np.eye(4)
cam_node = self.scene.add(camera, pose=camera_pose)
if self.wireframe:
render_flags = RenderFlags.RGBA | RenderFlags.ALL_WIREFRAME
else:
render_flags = RenderFlags.RGBA
rgb, _ = self.renderer.render(self.scene, flags=render_flags)
valid_mask = (rgb[:, :, -1] > 0)[:, :, np.newaxis]
output_img = rgb[:, :, :-1] * valid_mask + (1 - valid_mask) * img
image = output_img.astype(np.uint8)
self.scene.remove_node(mesh_node)
self.scene.remove_node(cam_node)
return image
def update_graph(self, scene_graph, materials_only):
"""Update scene graph.
This function rebuild scene completely each time when something changed.
TODO: update changed nodes instead.
Arguments:
scene_graph {SceneGraph} -- scene description
materials_only {bool} -- update only shape materials
"""
for uid, node in self._bullet_nodes.items():
self.remove_node(node)
self._bullet_nodes = {}
self._seg_node_map = {}
for uid, body in scene_graph.nodes.items():
node = pyr.Node(uid)
self.add_node(node)
self._bullet_nodes[uid] = node
for shape in body.shapes:
if shape.material is None:
material = None
else:
material = pyr.MetallicRoughnessMaterial(
baseColorFactor=shape.material.diffuse_color,
metallicFactor=0.2,
roughnessFactor=0.8,
alphaMode='BLEND')
texture = shape.material.diffuse_texture
if texture is not None:
if texture.bitmap is not None:
image = Image.fromarray(texture.bitmap)
else:
image = Image.open(
os.path.abspath(texture.filename))
texture = pyr.Texture(source=image,
source_channels=image.mode)
material.baseColorTexture = texture
if shape.mesh is None:
mesh = primitive_mesh(shape)
elif shape.mesh.data is None:
mesh = trimesh.load(os.path.abspath(shape.mesh.filename))
else:
data = shape.mesh.data
mesh = trimesh.Trimesh(
vertices=data.vertices,
vertex_normals=data.normals,
faces=data.faces,
visual=trimesh.visual.TextureVisuals(uv=data.uvs))
mesh = pyr.Mesh.from_trimesh(mesh, material=material)
mesh_node = self.add(mesh,
pose=shape.pose.matrix.T,
parent_node=node)
self._seg_node_map[mesh_node] = mask_to_rgb(
body.body, body.link)
def create_mesh_pyrender(vert, faces, col):
import trimesh
import pyrender
mesh = trimesh.Trimesh(vert, faces, process=False)
material = pyrender.MetallicRoughnessMaterial(metallicFactor=0.0,
alphaMode='OPAQUE',
baseColorFactor=col)
mesh = pyrender.Mesh.from_trimesh(mesh, material=material)
return mesh
def __call__(self, vertices, camera_translation, image):
material = pyrender.MetallicRoughnessMaterial(
metallicFactor=0.2,
alphaMode='OPAQUE',
baseColorFactor=(0.8, 0.3, 0.3, 1.0))
camera_translation[0] *= -1.
mesh = trimesh.Trimesh(vertices, self.faces)
rot = trimesh.transformations.rotation_matrix(np.radians(180),
[1, 0, 0])
mesh.apply_transform(rot)
#rot1 = np.eye(4)
#rot1[:3,:3] = camera_rotation.T
#mesh.apply_transform(rot1)
mesh = pyrender.Mesh.from_trimesh(mesh, material=material)
scene = pyrender.Scene(ambient_light=(0.5, 0.5, 0.5))
scene.add(mesh, 'mesh')
camera_pose = np.eye(4)
camera_pose[:3, 3] = camera_translation
#camera_pose[:3,:3] = camera_rotation
#print(camera_pose)
camera = pyrender.IntrinsicsCamera(fx=self.focal_length,
fy=self.focal_length,
cx=self.camera_center[0],
cy=self.camera_center[1])
#camera = pyrender.PerspectiveCamera(yfov=np.pi / 3.0, aspectRatio=1.0)
scene.add(camera, pose=camera_pose)
#light = pyrender.SpotLight(color=np.ones(3), intensity=3.0,
# innerConeAngle=np.pi/16.0,
# outerConeAngle=np.pi/6.0)
#scene.add(light, pose=camera_pose)
light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=1)
#scene.add(light, pose=camera_pose)
light_pose = np.eye(4)
#light_pose[:3,:3] = camera_rotation
light_pose[:3, 3] = np.array([0, -1, 1])
scene.add(light, pose=light_pose)
light_pose[:3, 3] = np.array([0, 1, 1])
scene.add(light, pose=light_pose)
light_pose[:3, 3] = np.array([1, 1, 2])
scene.add(light, pose=light_pose)
color, rend_depth = self.renderer.render(
scene, flags=pyrender.RenderFlags.RGBA)
color = color.astype(np.float32) / 255.0
valid_mask = (rend_depth > 0)[:, :, None]
output_img = (color[:, :, :3] * valid_mask + (1 - valid_mask) * image)
return output_img
def render(self, render_data, cameras, images,
use_white=False,
ret_depth=False, ret_color=False):
# Need to flip x-axis
rot = trimesh.transformations.rotation_matrix(
np.radians(180), [1, 0, 0])
output_images, output_colors, output_depths = [], [], []
for nv, img_ in enumerate(images):
if use_white:
img = np.zeros_like(img_, dtype=np.uint8) + 255
else:
img = img_.copy()
K, R, T = cameras['K'][nv], cameras['R'][nv], cameras['T'][nv]
self.renderer.viewport_height = img.shape[0]
self.renderer.viewport_width = img.shape[1]
scene = pyrender.Scene(bg_color=self.bg_color,
ambient_light=self.ambient_light)
for trackId, data in render_data.items():
vert = data['vertices'].copy()
faces = data['faces']
# 如果使用了vid这个键，那么可视化的颜色使用vid的颜色
col = get_colors(data.get('vid', trackId))
vert = vert @ R.T + T.T
mesh = trimesh.Trimesh(vert, faces)
mesh.apply_transform(rot)
material = pyrender.MetallicRoughnessMaterial(
metallicFactor=0.0,
alphaMode='OPAQUE',
baseColorFactor=col)
mesh = pyrender.Mesh.from_trimesh(
mesh,
material=material)
scene.add(mesh, 'mesh')
camera_pose = np.eye(4)
camera = pyrender.camera.IntrinsicsCamera(fx=K[0, 0], fy=K[1, 1], cx=K[0, 2], cy=K[1, 2])
scene.add(camera, pose=camera_pose)
self.add_light(scene)
# Alpha channel was not working previously need to check again
# Until this is fixed use hack with depth image to get the opacity
rend_rgba, rend_depth = self.renderer.render(scene, flags=flags)
if rend_rgba.shape[2] == 3: # fail to generate transparent channel
valid_mask = (rend_depth > 0)[:, :, None]
rend_rgba = np.dstack((rend_rgba, (valid_mask*255).astype(np.uint8)))
rend_cat = cv2.addWeighted(cv2.bitwise_and(img, 255 - rend_rgba[:, :, 3:4].repeat(3, 2)), 1, rend_rgba[:, :, :3], 1, 0)
output_colors.append(rend_rgba)
output_depths.append(rend_depth)
output_images.append(rend_cat)
if ret_depth:
return output_images, output_depths
elif ret_color:
return output_colors
else:
return output_images
def __call__(self,
verts,
cam_trans,
img=None,
angle=None,
axis=None,
mesh_filename=None,
color=[0.8, 0.3, 0.3],
return_mask=False):
mesh = trimesh.Trimesh(verts, self.faces, process=False)
Rx = trimesh.transformations.rotation_matrix(np.radians(180),
[1, 0, 0])
mesh.apply_transform(Rx)
cam_trans[0] *= -1.
if angle and axis:
# Apply given mesh rotation to the mesh - useful for rendering from different views
R = trimesh.transformations.rotation_matrix(
math.radians(angle), axis)
mesh.apply_transform(R)
material = pyrender.MetallicRoughnessMaterial(
metallicFactor=0.2,
alphaMode='OPAQUE',
baseColorFactor=(color[0], color[1], color[2], 1.0))
mesh = pyrender.Mesh.from_trimesh(mesh, material=material)
mesh_node = self.scene.add(mesh, 'mesh')
camera_pose = np.eye(4)
camera_pose[:3, 3] = cam_trans
camera = pyrender.IntrinsicsCamera(fx=config.FOCAL_LENGTH,
fy=config.FOCAL_LENGTH,
cx=self.camera_center[0],
cy=self.camera_center[1])
cam_node = self.scene.add(camera, pose=camera_pose)
rgb, rend_depth = self.renderer.render(self.scene,
flags=RenderFlags.RGBA)
valid_mask = (rend_depth > 0)
if return_mask:
return valid_mask
else:
if img is None:
img = np.zeros((self.img_res, self.img_res, 3))
valid_mask = valid_mask[:, :, None]
output_img = rgb[:, :, :-1] * valid_mask + (1 - valid_mask) * img
image = output_img.astype(np.uint8)
self.scene.remove_node(mesh_node)
self.scene.remove_node(cam_node)
return image
def update_mesh(self, mesh):
material = pyrender.MetallicRoughnessMaterial(metallicFactor=.0,
alphaMode='OPAQUE',
baseColorFactor=(1., 1.,
.9, 1.))
node = self.scene.get_nodes(name='mesh')
if len(node) > 1:
self.scene.remove_node(node.pop())
mesh = pyrender.Mesh.from_trimesh(mesh,
material=material,
smooth=False)
self.scene.add(mesh, 'mesh')
def render_smpl(vertices, faces, image, intrinsics, pose, transl,
alpha=1.0, filename='render_sample.png'):
img_size = image.shape[-2]
material = pyrender.MetallicRoughnessMaterial(
metallicFactor=0.2,
alphaMode='OPAQUE',
baseColorFactor=(0.8, 0.3, 0.3, 1.0))
# Generate SMPL vertices mesh
mesh = trimesh.Trimesh(vertices, faces)
# Default rotation of SMPL body model
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(ambient_light=(0.5, 0.5, 0.5))
scene.add(mesh, 'mesh')
camera_pose = np.eye(4)
camera_pose[:3, :3] = pose
camera_pose[:3, 3] = transl
camera = pyrender.IntrinsicsCamera(fx=intrinsics[0, 0], fy=intrinsics[1, 1],
cx=intrinsics[0, 2], cy=intrinsics[1, 2])
scene.add(camera, pose=camera_pose)
# Light information
light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=1)
light_pose = np.eye(4)
light_pose[:3, 3] = np.array([0, -1, 1])
scene.add(light, pose=light_pose)
light_pose[:3, 3] = np.array([0, 1, 1])
scene.add(light, pose=light_pose)
light_pose[:3, 3] = np.array([1, 1, 2])
scene.add(light, pose=light_pose)
renderer = pyrender.OffscreenRenderer(
viewport_width=img_size, viewport_height=img_size, point_size=1.0)
color, rend_depth = renderer.render(scene, flags=pyrender.RenderFlags.RGBA)
valid_mask = (rend_depth > 0)[:,:,None]
color = color.astype(np.float32) / 255.0
valid_mask = (rend_depth > 0)[:,:,None]
output_img = color[:, :, :3] * valid_mask + (1 - valid_mask) * image / 255.0
cv2.imwrite(filename, 255 * output_img)
def __call__(self, vertices, camera_translation, image, angle=180, rot_axis=[1, 0, 0]):
material = pyrender.MetallicRoughnessMaterial(
metallicFactor=0.2,
alphaMode='OPAQUE',
baseColorFactor=(0.8, 0.3, 0.3, 1.0))
camera_translation[0] *= -1.
mesh = trimesh.Trimesh(vertices, self.faces)
if angle != 180:
rot = trimesh.transformations.rotation_matrix(
np.radians(180), [1, 0, 0])
mesh.apply_transform(rot)
rot = trimesh.transformations.rotation_matrix(
np.radians(90), [0, 1, 0])
mesh.apply_transform(rot)
else:
rot = trimesh.transformations.rotation_matrix(
np.radians(angle), rot_axis)
mesh.apply_transform(rot)
mesh = pyrender.Mesh.from_trimesh(mesh, material=material)
scene = pyrender.Scene(ambient_light=(0.5, 0.5, 0.5))
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=self.camera_center[0], cy=self.camera_center[1])
scene.add(camera, pose=camera_pose)
light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=1)
light_pose = np.eye(4)
light_pose[:3, 3] = np.array([0, -1, 1])
scene.add(light, pose=light_pose)
light_pose[:3, 3] = np.array([0, 1, 1])
scene.add(light, pose=light_pose)
light_pose[:3, 3] = np.array([1, 1, 2])
scene.add(light, pose=light_pose)
color, rend_depth = self.renderer.render(scene, flags=pyrender.RenderFlags.RGBA)
color = color.astype(np.float32) / 255.0
valid_mask = (rend_depth > 0)[:,:,None]
output_img = (color[:, :, :3] * valid_mask +
(1 - valid_mask) * image)
return output_img
def __init__(self):
self.scene = pyrender.Scene(ambient_light=[0.3, 0.3, 0.3, 0.],
bg_color=[1.0, 1.0, 1.0, 0.])
self.viewer = pyrender.Viewer(self.scene,
viewport_size=(1400, 900),
run_in_thread=True,
viewer_flags=viewer_flags,
use_raymond_lighting=True)
self.mesh_node = None
self.material = pyrender.MetallicRoughnessMaterial(
metallicFactor=0.5,
roughnessFactor=0.8,
alphaMode='OPAQUE',
baseColorFactor=(.6, .6, .6, 1.))
def create_mesh(self, vertices, faces, vertex_color=(0.9, 0.9, 0.7, 1.0)):
tri_mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
rot = trimesh.transformations.rotation_matrix(np.radians(180),
[1, 0, 0])
tri_mesh.apply_transform(rot)
meshes = []
material = pyrender.MetallicRoughnessMaterial(
metallicFactor=0.0, baseColorFactor=vertex_color)
mesh = pyrender.Mesh.from_trimesh(tri_mesh, material=material)
meshes.append(mesh)
return meshes
def __call__(self,
vertices,
camera_pose,
image,
color=(0.8, 0.3, 0.3, 1.0)):
material = pyrender.MetallicRoughnessMaterial(metallicFactor=0.2,
alphaMode='OPAQUE',
baseColorFactor=color)
mesh = trimesh.Trimesh(vertices, self.faces)
# Rotate mesh 180 deg around x (pyrender coordinate frame)
rot = trimesh.transformations.rotation_matrix(np.radians(180),
[1, 0, 0])
mesh.apply_transform(rot)
mesh = pyrender.Mesh.from_trimesh(mesh, material=material)
# Rotate trafo 180 deg around x (pyrender coordinate frame)
Rx = np.array(
[[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]],
dtype=float)
camera_pose = np.dot(camera_pose, Rx)
scene = pyrender.Scene(ambient_light=(0.5, 0.5, 0.5))
scene.add(mesh, 'mesh')
camera = pyrender.IntrinsicsCamera(fx=self.focal_length[0],
fy=self.focal_length[1],
cx=self.camera_center[0],
cy=self.camera_center[1])
scene.add(camera, pose=camera_pose)
light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=1)
light_pose = np.eye(4)
light_pose[:3, 3] = np.array([0, -1, 1])
scene.add(light, pose=light_pose)
light_pose[:3, 3] = np.array([0, 1, 1])
scene.add(light, pose=light_pose)
light_pose[:3, 3] = np.array([1, 1, 2])
scene.add(light, pose=light_pose)
color, rend_depth = self.renderer.render(
scene, flags=pyrender.RenderFlags.RGBA)
valid_mask = (rend_depth > 0)[:, :, None]
output_img = (color[:, :, :3] * valid_mask +
(1 - valid_mask) * image).astype(np.uint8)
return output_img
def create_mesh(vertices, faces, color=(0.3, 0.3, 0.3, 1.0),
wireframe=False):
tri_mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
rot = trimesh.transformations.rotation_matrix(np.radians(180),
[1, 0, 0])
tri_mesh.apply_transform(rot)
material = pyrender.MetallicRoughnessMaterial(
metallicFactor=0.0,
alphaMode='BLEND',
baseColorFactor=color)
return pyrender.Mesh.from_trimesh(
tri_mesh,
material=material)
def render_mesh(mesh_trimesh, camera_center, camera_transl, focal_length, img_width, img_height):
import pyrender
import trimesh
material = pyrender.MetallicRoughnessMaterial(
metallicFactor=0.0,
alphaMode='OPAQUE',
baseColorFactor=(1.0, 1.0, 0.9, 1.0))
script_dir = os.path.dirname(os.path.realpath(__file__))
vertex_colors = np.loadtxt(os.path.join(script_dir, 'smplx_verts_colors.txt'))
mesh_new = trimesh.Trimesh(vertices=mesh_trimesh.vertices, faces=mesh_trimesh.faces, vertex_colors=vertex_colors)
mesh_new.vertex_colors = vertex_colors
print("mesh visual kind: %s" % mesh_new.visual.kind)
#mesh = pyrender.Mesh.from_points(out_mesh.vertices, colors=vertex_colors)
mesh = pyrender.Mesh.from_trimesh(mesh_new, smooth=False, wireframe=False)
scene = pyrender.Scene(bg_color=[1.0, 1.0, 1.0, 0.0],
ambient_light=(0.3, 0.3, 0.3))
#scene = pyrender.Scene(bg_color=[0.0, 0.0, 0.0, 0.0])
scene.add(mesh, 'mesh')
camera_pose = np.eye(4)
camera_pose[:3, 3] = camera_transl
camera = pyrender.camera.IntrinsicsCamera(
fx=focal_length, fy=focal_length,
cx=camera_center[0], cy=camera_center[1])
scene.add(camera, pose=camera_pose)
light = pyrender.light.DirectionalLight()
scene.add(light)
r = pyrender.OffscreenRenderer(viewport_width=img_width,
viewport_height=img_height,
point_size=1.0)
color, _ = r.render(scene, flags=pyrender.RenderFlags.RGBA)
color = color.astype(np.float32) / 255.0
output_img = color[:, :, 0:3]
output_img = (output_img * 255).astype(np.uint8)
return output_img
def deodr_mesh_to_pyrender(deodr_mesh):
# trimesh and pyrender do to handle faces indices for texture
# that are different from face indices for the 3d vertices
# we need to duplicate vertices
faces, mask_v, mask_vt = trimesh.visual.texture.unmerge_faces(
deodr_mesh.faces, deodr_mesh.faces_uv)
vertices = deodr_mesh.vertices[mask_v]
deodr_mesh.compute_vertex_normals()
vertex_normals = deodr_mesh.vertex_normals[mask_v]
uv = deodr_mesh.uv[mask_vt]
pyrender_uv = np.column_stack((
(uv[:, 0]) / deodr_mesh.texture.shape[0],
(1 - uv[:, 1] / deodr_mesh.texture.shape[1]),
))
material = None
poses = None
color_0 = None
if material is None:
base_color_texture = pyrender.texture.Texture(
source=deodr_mesh.texture, source_channels="RGB")
material = pyrender.MetallicRoughnessMaterial(
alphaMode="BLEND",
baseColorFactor=[1, 1, 1, 1.0],
metallicFactor=0,
roughnessFactor=1,
baseColorTexture=base_color_texture,
)
material.wireframe = False
primitive = pyrender.Primitive(
positions=vertices,
normals=vertex_normals,
texcoord_0=pyrender_uv,
color_0=color_0,
indices=faces,
material=material,
mode=pyrender.constants.GLTF.TRIANGLES,
poses=poses,
)
return pyrender.Mesh(primitives=[primitive])
def render_mesh(img, mesh, face, cam_param):
# mesh
mesh = trimesh.Trimesh(mesh, face)
rot = trimesh.transformations.rotation_matrix(np.radians(180), [1, 0, 0])
mesh.apply_transform(rot)
material = pyrender.MetallicRoughnessMaterial(metallicFactor=0.0,
alphaMode='OPAQUE',
baseColorFactor=(1.0, 1.0,
0.9, 1.0))
mesh = pyrender.Mesh.from_trimesh(mesh, material=material, smooth=False)
scene = pyrender.Scene(ambient_light=(0.3, 0.3, 0.3))
scene.add(mesh, 'mesh')
focal, princpt = cam_param['focal'], cam_param['princpt']
camera = pyrender.IntrinsicsCamera(fx=focal[0],
fy=focal[1],
cx=princpt[0],
cy=princpt[1])
scene.add(camera)
# renderer
renderer = pyrender.OffscreenRenderer(viewport_width=img.shape[1],
viewport_height=img.shape[0],
point_size=1.0)
# light
light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=0.8)
light_pose = np.eye(4)
light_pose[:3, 3] = np.array([0, -1, 1])
scene.add(light, pose=light_pose)
light_pose[:3, 3] = np.array([0, 1, 1])
scene.add(light, pose=light_pose)
light_pose[:3, 3] = np.array([1, 1, 2])
scene.add(light, pose=light_pose)
# render
rgb, depth = renderer.render(scene, flags=pyrender.RenderFlags.RGBA)
rgb = rgb[:, :, :3].astype(np.float32)
valid_mask = (depth > 0)[:, :, None]
# save to image
img = rgb * valid_mask + img * (1 - valid_mask)
return img
def __call__(self, vertices, faces):
material = pyrender.MetallicRoughnessMaterial(
metallicFactor=0.2,
alphaMode='OPAQUE',
baseColorFactor=(0.8, 0.3, 0.3, 1.0))
camera_translation = np.array([0.0, 0.0, 50.0])
mesh = trimesh.Trimesh(vertices[0], faces[0], process=False)
# 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(ambient_light=(0.5, 0.5, 0.5))
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=self.camera_center[0],
cy=self.camera_center[1])
scene.add(camera, pose=camera_pose)
light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=1)
light_pose = np.eye(4)
light_pose[:3, 3] = np.array([0, -1, 1])
scene.add(light, pose=light_pose)
light_pose[:3, 3] = np.array([0, 1, 1])
scene.add(light, pose=light_pose)
light_pose[:3, 3] = np.array([1, 1, 2])
scene.add(light, pose=light_pose)
color, rend_depth = self.renderer.render(
scene, flags=pyrender.RenderFlags.RGBA)
color = color.astype(np.float32) / 255.0
return torch.from_numpy(color).float().unsqueeze(0)
def main():
dataset = get_dataset('s0_train')
# Load pre-generated grasps for YCB objects.
ycb_grasp_file = os.path.join(os.path.dirname(__file__), "..", "assets",
"ycb_farthest_100_grasps.json")
with open(ycb_grasp_file, 'r') as f:
ycb_grasps = json.load(f)
# Load simplified panda gripper mesh.
gripper_mesh_file = os.path.join(os.path.dirname(__file__), "..", "assets",
"panda_tubes.obj")
gripper_mesh = trimesh.load(gripper_mesh_file)
gripper_material = pyrender.MetallicRoughnessMaterial(
alphaMode="BLEND",
doubleSided=True,
baseColorFactor=(0.00, 1.00, 0.04, 1.00),
metallicFactor=0.0)
# Visualize pre-generated grasps for each YCB object.
for ycb_id, name in dataset.ycb_classes.items():
if name not in ycb_grasps.keys():
print('{} does not have pre-generated grasps: skip.'.format(name))
continue
scene = pyrender.Scene(ambient_light=np.array([0.5, 0.5, 0.5, 1.0]))
obj_mesh = trimesh.load(dataset.obj_file[ycb_id])
scene.add(pyrender.Mesh.from_trimesh(obj_mesh))
for grasp in ycb_grasps[name]:
grasp_mesh = copy.deepcopy(gripper_mesh).apply_transform(grasp),
scene.add(pyrender.Mesh.from_trimesh(grasp_mesh, gripper_material))
print('Visualizing pre-generated grasps of {}'.format(name))
print('Close the window to visualize next object.')
pyrender.Viewer(scene, use_raymond_lighting=True)
def imshow_mesh_3d(img,
vertices,
faces,
camera_center,
focal_length,
colors=(76, 76, 204)):
"""Render 3D meshes on background image.
Args:
img(np.ndarray): Background image.
vertices (list of np.ndarray): Vetrex coordinates in camera space.
faces (list of np.ndarray): Faces of meshes.
camera_center ([2]): Center pixel.
focal_length ([2]): Focal length of camera.
colors (list[str or tuple or Color]): A list of mesh colors.
"""
H, W, C = img.shape
if not has_pyrender:
warnings.warn('pyrender package is not installed.')
return img
if not has_trimesh:
warnings.warn('trimesh package is not installed.')
return img
try:
renderer = pyrender.OffscreenRenderer(viewport_width=W,
viewport_height=H)
except (ImportError, RuntimeError):
warnings.warn('pyrender package is not installed correctly.')
return img
if not isinstance(colors, list):
colors = [colors for _ in range(len(vertices))]
colors = [color_val(c) for c in colors]
depth_map = np.ones([H, W]) * np.inf
output_img = img
for idx in range(len(vertices)):
color = colors[idx]
color = [c / 255.0 for c in color]
color.append(1.0)
vert = vertices[idx]
face = faces[idx]
material = pyrender.MetallicRoughnessMaterial(metallicFactor=0.2,
alphaMode='OPAQUE',
baseColorFactor=color)
mesh = trimesh.Trimesh(vert, face)
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(ambient_light=(0.5, 0.5, 0.5))
scene.add(mesh, 'mesh')
camera_pose = np.eye(4)
camera = pyrender.IntrinsicsCamera(fx=focal_length[0],
fy=focal_length[1],
cx=camera_center[0],
cy=camera_center[1],
zfar=1e5)
scene.add(camera, pose=camera_pose)
light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=1)
light_pose = np.eye(4)
light_pose[:3, 3] = np.array([0, -1, 1])
scene.add(light, pose=light_pose)
light_pose[:3, 3] = np.array([0, 1, 1])
scene.add(light, pose=light_pose)
light_pose[:3, 3] = np.array([1, 1, 2])
scene.add(light, pose=light_pose)
color, rend_depth = renderer.render(scene,
flags=pyrender.RenderFlags.RGBA)
valid_mask = (rend_depth < depth_map) * (rend_depth > 0)
depth_map[valid_mask] = rend_depth[valid_mask]
valid_mask = valid_mask[:, :, None]
output_img = (valid_mask * color[:, :, :3] +
(1 - valid_mask) * output_img)
return output_img
# mask = np.where(referenced==True)[0]
# stacked = [input_mesh.vertices * (10 ** digits_vertex)]
# stacked = np.column_stack(stacked).round().astype(np.int64)
# u, i = unique_rows(stacked[referenced])
# inverse = np.zeros(len(input_mesh.vertices), dtype=np.int64)
# inverse[referenced] = i
# mask = np.nonzero(referenced)[0][u]
# collisions_mesh = trimesh.Trimesh(input_mesh.vertices, collisions_faces)
# trimesh.smoothing.filter_humphrey(collisions_mesh, iterations=20)
# input_mesh.vertices[mask] = collisions_mesh.vertices
material = pyrender.MetallicRoughnessMaterial(
metallicFactor=0.0,
alphaMode='BLEND',
baseColorFactor=[0.3, 0.3, 0.3, 0.99])
recv_material = pyrender.MetallicRoughnessMaterial(
metallicFactor=0.0,
alphaMode='BLEND',
baseColorFactor=[0.0, 0.9, 0.0, 1.0])
intr_material = pyrender.MetallicRoughnessMaterial(
metallicFactor=0.0,
alphaMode='BLEND',
baseColorFactor=[0.9, 0.0, 0.0, 1.0])
main_mesh = pyrender.Mesh.from_trimesh(input_mesh, material=material)
recv_mesh = pyrender.Mesh.from_trimesh(r_mesh, material=recv_material)
intr_mesh = pyrender.Mesh.from_trimesh(i_mesh, material=intr_material)
est_params = {}
for key, val in data.items():
if key == 'body_pose' and use_vposer:
body_pose = vposer.decode(
pose_embedding, output_type='aa').view(1, -1)
if model_type == 'smpl':
wrist_pose = torch.zeros([body_pose.shape[0], 6],
dtype=body_pose.dtype,
device=body_pose.device)
body_pose = torch.cat([body_pose, wrist_pose], dim=1)
est_params['body_pose'] = body_pose
else:
est_params[key] = torch.tensor(val, dtype=dtype, device=device)
model_output = model(**est_params)
vertices = model_output.vertices.detach().cpu().numpy().squeeze()
out_mesh = trimesh.Trimesh(vertices, model.faces, process=False)
material = pyrender.MetallicRoughnessMaterial(
metallicFactor=0.0,
alphaMode='OPAQUE',
baseColorFactor=(1.0, 1.0, 0.9, 1.0))
mesh = pyrender.Mesh.from_trimesh(
out_mesh,
material=material)
scene = pyrender.Scene(bg_color=[0.0, 0.0, 0.0, 0.0],
ambient_light=(0.3, 0.3, 0.3))
scene.add(mesh, 'mesh')
pyrender.Viewer(scene, use_raymond_lighting=True)
def __call__(self, images, vertices, translation):
# List to store rendered scenes
output_images = []
# Need to flip x-axis
rot = trimesh.transformations.rotation_matrix(np.radians(180),
[1, 0, 0])
# For all iamges
for i in range(len(images)):
img = images[i].cpu().numpy().transpose(1, 2, 0)
self.renderer.viewport_height = img.shape[0]
self.renderer.viewport_width = img.shape[1]
verts = vertices[i].detach().cpu().numpy()
mesh_trans = translation[i].cpu().numpy()
verts = verts + mesh_trans[:, None, ]
num_people = verts.shape[0]
# Create a scene for each image and render all meshes
scene = pyrender.Scene(bg_color=[0.0, 0.0, 0.0, 0.0],
ambient_light=(0.5, 0.5, 0.5))
# Create camera. Camera will always be at [0,0,0]
# CHECK If I need to swap x and y
camera_center = np.array([img.shape[1] / 2., img.shape[0] / 2.])
camera_pose = np.eye(4)
camera = pyrender.camera.IntrinsicsCamera(fx=self.focal_length,
fy=self.focal_length,
cx=camera_center[0],
cy=camera_center[1])
scene.add(camera, pose=camera_pose)
# Create light source
light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0],
intensity=1)
# for every person in the scene
for n in range(num_people):
mesh = trimesh.Trimesh(verts[n], self.faces)
mesh.apply_transform(rot)
trans = 0 * mesh_trans[n]
trans[0] *= -1
trans[2] *= -1
material = pyrender.MetallicRoughnessMaterial(
metallicFactor=0.2,
alphaMode='OPAQUE',
baseColorFactor=colors[n % len(colors)])
mesh = pyrender.Mesh.from_trimesh(mesh, material=material)
scene.add(mesh, 'mesh')
# Use 3 directional lights
light_pose = np.eye(4)
light_pose[:3, 3] = np.array([0, -1, 1]) + trans
scene.add(light, pose=light_pose)
light_pose[:3, 3] = np.array([0, 1, 1]) + trans
scene.add(light, pose=light_pose)
light_pose[:3, 3] = np.array([1, 1, 2]) + trans
scene.add(light, pose=light_pose)
# Alpha channel was not working previously need to check again
# Until this is fixed use hack with depth image to get the opacity
color, rend_depth = self.renderer.render(
scene, flags=pyrender.RenderFlags.RGBA)
# color = color[::-1,::-1]
# rend_depth = rend_depth[::-1,::-1]
color = color.astype(np.float32) / 255.0
valid_mask = (rend_depth > 0)[:, :, None]
output_img = (color[:, :, :3] * valid_mask +
(1 - valid_mask) * img)
output_img = np.transpose(output_img, (2, 0, 1))
output_images.append(output_img)
return output_images
def fit_single_frame(
img,
keypoints,
init_trans,
scan,
scene_name,
body_model,
camera,
joint_weights,
body_pose_prior,
jaw_prior,
left_hand_prior,
right_hand_prior,
shape_prior,
expr_prior,
angle_prior,
result_fn='out.pkl',
mesh_fn='out.obj',
body_scene_rendering_fn='body_scene.png',
out_img_fn='overlay.png',
loss_type='smplify',
use_cuda=True,
init_joints_idxs=(9, 12, 2, 5),
use_face=True,
use_hands=True,
data_weights=None,
body_pose_prior_weights=None,
hand_pose_prior_weights=None,
jaw_pose_prior_weights=None,
shape_weights=None,
expr_weights=None,
hand_joints_weights=None,
face_joints_weights=None,
depth_loss_weight=1e2,
interpenetration=True,
coll_loss_weights=None,
df_cone_height=0.5,
penalize_outside=True,
max_collisions=8,
point2plane=False,
part_segm_fn='',
focal_length_x=5000.,
focal_length_y=5000.,
side_view_thsh=25.,
rho=100,
vposer_latent_dim=32,
vposer_ckpt='',
use_joints_conf=False,
interactive=True,
visualize=False,
save_meshes=True,
degrees=None,
batch_size=1,
dtype=torch.float32,
ign_part_pairs=None,
left_shoulder_idx=2,
right_shoulder_idx=5,
####################
### PROX
render_results=True,
camera_mode='moving',
## Depth
s2m=False,
s2m_weights=None,
m2s=False,
m2s_weights=None,
rho_s2m=1,
rho_m2s=1,
init_mode=None,
trans_opt_stages=None,
viz_mode='mv',
#penetration
sdf_penetration=False,
sdf_penetration_weights=0.0,
sdf_dir=None,
cam2world_dir=None,
#contact
contact=False,
rho_contact=1.0,
contact_loss_weights=None,
contact_angle=15,
contact_body_parts=None,
body_segments_dir=None,
load_scene=False,
scene_dir=None,
**kwargs):
assert batch_size == 1, 'PyTorch L-BFGS only supports batch_size == 1'
body_model.reset_params()
body_model.transl.requires_grad = True
device = torch.device('cuda') if use_cuda else torch.device('cpu')
if visualize:
pil_img.fromarray((img * 255).astype(np.uint8)).show()
if degrees is None:
degrees = [0, 90, 180, 270]
if data_weights is None:
data_weights = [
1,
] * 5
if body_pose_prior_weights is None:
body_pose_prior_weights = [4.04 * 1e2, 4.04 * 1e2, 57.4, 4.78]
msg = ('Number of Body pose prior weights {}'.format(
len(body_pose_prior_weights)) +
' does not match the number of data term weights {}'.format(
len(data_weights)))
assert (len(data_weights) == len(body_pose_prior_weights)), msg
if use_hands:
if hand_pose_prior_weights is None:
hand_pose_prior_weights = [1e2, 5 * 1e1, 1e1, .5 * 1e1]
msg = ('Number of Body pose prior weights does not match the' +
' number of hand pose prior weights')
assert (
len(hand_pose_prior_weights) == len(body_pose_prior_weights)), msg
if hand_joints_weights is None:
hand_joints_weights = [0.0, 0.0, 0.0, 1.0]
msg = ('Number of Body pose prior weights does not match the' +
' number of hand joint distance weights')
assert (
len(hand_joints_weights) == len(body_pose_prior_weights)), msg
if shape_weights is None:
shape_weights = [1e2, 5 * 1e1, 1e1, .5 * 1e1]
msg = ('Number of Body pose prior weights = {} does not match the' +
' number of Shape prior weights = {}')
assert (len(shape_weights) == len(body_pose_prior_weights)), msg.format(
len(shape_weights), len(body_pose_prior_weights))
if use_face:
if jaw_pose_prior_weights is None:
jaw_pose_prior_weights = [[x] * 3 for x in shape_weights]
else:
jaw_pose_prior_weights = map(lambda x: map(float, x.split(',')),
jaw_pose_prior_weights)
jaw_pose_prior_weights = [list(w) for w in jaw_pose_prior_weights]
msg = ('Number of Body pose prior weights does not match the' +
' number of jaw pose prior weights')
assert (
len(jaw_pose_prior_weights) == len(body_pose_prior_weights)), msg
if expr_weights is None:
expr_weights = [1e2, 5 * 1e1, 1e1, .5 * 1e1]
msg = ('Number of Body pose prior weights = {} does not match the' +
' number of Expression prior weights = {}')
assert (len(expr_weights) == len(body_pose_prior_weights)), msg.format(
len(body_pose_prior_weights), len(expr_weights))
if face_joints_weights is None:
face_joints_weights = [0.0, 0.0, 0.0, 1.0]
msg = ('Number of Body pose prior weights does not match the' +
' number of face joint distance weights')
assert (len(face_joints_weights) == len(body_pose_prior_weights)), msg
if coll_loss_weights is None:
coll_loss_weights = [0.0] * len(body_pose_prior_weights)
msg = ('Number of Body pose prior weights does not match the' +
' number of collision loss weights')
assert (len(coll_loss_weights) == len(body_pose_prior_weights)), msg
use_vposer = kwargs.get('use_vposer', True)
vposer, pose_embedding = [
None,
] * 2
if use_vposer:
pose_embedding = torch.zeros([batch_size, 32],
dtype=dtype,
device=device,
requires_grad=True)
vposer_ckpt = osp.expandvars(vposer_ckpt)
vposer, _ = load_vposer(vposer_ckpt, vp_model='snapshot')
vposer = vposer.to(device=device)
vposer.eval()
if use_vposer:
body_mean_pose = torch.zeros([batch_size, vposer_latent_dim],
dtype=dtype)
else:
body_mean_pose = body_pose_prior.get_mean().detach().cpu()
keypoint_data = torch.tensor(keypoints, dtype=dtype)
gt_joints = keypoint_data[:, :, :2]
if use_joints_conf:
joints_conf = keypoint_data[:, :, 2].reshape(1, -1)
# Transfer the data to the correct device
gt_joints = gt_joints.to(device=device, dtype=dtype)
if use_joints_conf:
joints_conf = joints_conf.to(device=device, dtype=dtype)
scan_tensor = None
if scan is not None:
scan_tensor = torch.tensor(scan.get('points'),
device=device,
dtype=dtype).unsqueeze(0)
# load pre-computed signed distance field
sdf = None
sdf_normals = None
grid_min = None
grid_max = None
voxel_size = None
if sdf_penetration:
with open(osp.join(sdf_dir, scene_name + '.json'), 'r') as f:
sdf_data = json.load(f)
grid_min = torch.tensor(np.array(sdf_data['min']),
dtype=dtype,
device=device)
grid_max = torch.tensor(np.array(sdf_data['max']),
dtype=dtype,
device=device)
grid_dim = sdf_data['dim']
voxel_size = (grid_max - grid_min) / grid_dim
sdf = np.load(osp.join(sdf_dir, scene_name + '_sdf.npy')).reshape(
grid_dim, grid_dim, grid_dim)
sdf = torch.tensor(sdf, dtype=dtype, device=device)
if osp.exists(osp.join(sdf_dir, scene_name + '_normals.npy')):
sdf_normals = np.load(
osp.join(sdf_dir, scene_name + '_normals.npy')).reshape(
grid_dim, grid_dim, grid_dim, 3)
sdf_normals = torch.tensor(sdf_normals, dtype=dtype, device=device)
else:
print("Normals not found...")
with open(os.path.join(cam2world_dir, scene_name + '.json'), 'r') as f:
cam2world = np.array(json.load(f))
R = torch.tensor(cam2world[:3, :3].reshape(3, 3),
dtype=dtype,
device=device)
t = torch.tensor(cam2world[:3, 3].reshape(1, 3),
dtype=dtype,
device=device)
# Create the search tree
search_tree = None
pen_distance = None
filter_faces = None
if interpenetration:
from mesh_intersection.bvh_search_tree import BVH
import mesh_intersection.loss as collisions_loss
from mesh_intersection.filter_faces import FilterFaces
assert use_cuda, 'Interpenetration term can only be used with CUDA'
assert torch.cuda.is_available(), \
'No CUDA Device! Interpenetration term can only be used' + \
' with CUDA'
search_tree = BVH(max_collisions=max_collisions)
pen_distance = \
collisions_loss.DistanceFieldPenetrationLoss(
sigma=df_cone_height, point2plane=point2plane,
vectorized=True, penalize_outside=penalize_outside)
if part_segm_fn:
# Read the part segmentation
part_segm_fn = os.path.expandvars(part_segm_fn)
with open(part_segm_fn, 'rb') as faces_parents_file:
face_segm_data = pickle.load(faces_parents_file,
encoding='latin1')
faces_segm = face_segm_data['segm']
faces_parents = face_segm_data['parents']
# Create the module used to filter invalid collision pairs
filter_faces = FilterFaces(
faces_segm=faces_segm,
faces_parents=faces_parents,
ign_part_pairs=ign_part_pairs).to(device=device)
# load vertix ids of contact parts
contact_verts_ids = ftov = None
if contact:
contact_verts_ids = []
for part in contact_body_parts:
with open(os.path.join(body_segments_dir, part + '.json'),
'r') as f:
data = json.load(f)
contact_verts_ids.append(list(set(data["verts_ind"])))
contact_verts_ids = np.concatenate(contact_verts_ids)
vertices = body_model(return_verts=True,
body_pose=torch.zeros((batch_size, 63),
dtype=dtype,
device=device)).vertices
vertices_np = vertices.detach().cpu().numpy().squeeze()
body_faces_np = body_model.faces_tensor.detach().cpu().numpy().reshape(
-1, 3)
m = Mesh(v=vertices_np, f=body_faces_np)
ftov = m.faces_by_vertex(as_sparse_matrix=True)
ftov = sparse.coo_matrix(ftov)
indices = torch.LongTensor(np.vstack((ftov.row, ftov.col))).to(device)
values = torch.FloatTensor(ftov.data).to(device)
shape = ftov.shape
ftov = torch.sparse.FloatTensor(indices, values, torch.Size(shape))
# Read the scene scan if any
scene_v = scene_vn = scene_f = None
if scene_name is not None:
if load_scene:
scene = Mesh(filename=os.path.join(scene_dir, scene_name + '.ply'))
scene.vn = scene.estimate_vertex_normals()
scene_v = torch.tensor(scene.v[np.newaxis, :],
dtype=dtype,
device=device).contiguous()
scene_vn = torch.tensor(scene.vn[np.newaxis, :],
dtype=dtype,
device=device)
scene_f = torch.tensor(scene.f.astype(int)[np.newaxis, :],
dtype=torch.long,
device=device)
# Weights used for the pose prior and the shape prior
opt_weights_dict = {
'data_weight': data_weights,
'body_pose_weight': body_pose_prior_weights,
'shape_weight': shape_weights
}
if use_face:
opt_weights_dict['face_weight'] = face_joints_weights
opt_weights_dict['expr_prior_weight'] = expr_weights
opt_weights_dict['jaw_prior_weight'] = jaw_pose_prior_weights
if use_hands:
opt_weights_dict['hand_weight'] = hand_joints_weights
opt_weights_dict['hand_prior_weight'] = hand_pose_prior_weights
if interpenetration:
opt_weights_dict['coll_loss_weight'] = coll_loss_weights
if s2m:
opt_weights_dict['s2m_weight'] = s2m_weights
if m2s:
opt_weights_dict['m2s_weight'] = m2s_weights
if sdf_penetration:
opt_weights_dict['sdf_penetration_weight'] = sdf_penetration_weights
if contact:
opt_weights_dict['contact_loss_weight'] = contact_loss_weights
keys = opt_weights_dict.keys()
opt_weights = [
dict(zip(keys, vals))
for vals in zip(*(opt_weights_dict[k] for k in keys
if opt_weights_dict[k] is not None))
]
for weight_list in opt_weights:
for key in weight_list:
weight_list[key] = torch.tensor(weight_list[key],
device=device,
dtype=dtype)
# load indices of the head of smpl-x model
with open(osp.join(body_segments_dir, 'body_mask.json'), 'r') as fp:
head_indx = np.array(json.load(fp))
N = body_model.get_num_verts()
body_indx = np.setdiff1d(np.arange(N), head_indx)
head_mask = np.in1d(np.arange(N), head_indx)
body_mask = np.in1d(np.arange(N), body_indx)
# The indices of the joints used for the initialization of the camera
init_joints_idxs = torch.tensor(init_joints_idxs, device=device)
edge_indices = kwargs.get('body_tri_idxs')
# which initialization mode to choose: similar traingles, mean of the scan or the average of both
if init_mode == 'scan':
init_t = init_trans
elif init_mode == 'both':
init_t = (init_trans.to(device) + fitting.guess_init(
body_model,
gt_joints,
edge_indices,
use_vposer=use_vposer,
vposer=vposer,
pose_embedding=pose_embedding,
model_type=kwargs.get('model_type', 'smpl'),
focal_length=focal_length_x,
dtype=dtype)) / 2.0
else:
init_t = fitting.guess_init(body_model,
gt_joints,
edge_indices,
use_vposer=use_vposer,
vposer=vposer,
pose_embedding=pose_embedding,
model_type=kwargs.get(
'model_type', 'smpl'),
focal_length=focal_length_x,
dtype=dtype)
camera_loss = fitting.create_loss('camera_init',
trans_estimation=init_t,
init_joints_idxs=init_joints_idxs,
depth_loss_weight=depth_loss_weight,
camera_mode=camera_mode,
dtype=dtype).to(device=device)
camera_loss.trans_estimation[:] = init_t
loss = fitting.create_loss(loss_type=loss_type,
joint_weights=joint_weights,
rho=rho,
use_joints_conf=use_joints_conf,
use_face=use_face,
use_hands=use_hands,
vposer=vposer,
pose_embedding=pose_embedding,
body_pose_prior=body_pose_prior,
shape_prior=shape_prior,
angle_prior=angle_prior,
expr_prior=expr_prior,
left_hand_prior=left_hand_prior,
right_hand_prior=right_hand_prior,
jaw_prior=jaw_prior,
interpenetration=interpenetration,
pen_distance=pen_distance,
search_tree=search_tree,
tri_filtering_module=filter_faces,
s2m=s2m,
m2s=m2s,
rho_s2m=rho_s2m,
rho_m2s=rho_m2s,
head_mask=head_mask,
body_mask=body_mask,
sdf_penetration=sdf_penetration,
voxel_size=voxel_size,
grid_min=grid_min,
grid_max=grid_max,
sdf=sdf,
sdf_normals=sdf_normals,
R=R,
t=t,
contact=contact,
contact_verts_ids=contact_verts_ids,
rho_contact=rho_contact,
contact_angle=contact_angle,
dtype=dtype,
**kwargs)
loss = loss.to(device=device)
with fitting.FittingMonitor(batch_size=batch_size,
visualize=visualize,
viz_mode=viz_mode,
**kwargs) as monitor:
img = torch.tensor(img, dtype=dtype)
H, W, _ = img.shape
# Reset the parameters to estimate the initial translation of the
# body model
if camera_mode == 'moving':
body_model.reset_params(body_pose=body_mean_pose)
# Update the value of the translation of the camera as well as
# the image center.
with torch.no_grad():
camera.translation[:] = init_t.view_as(camera.translation)
camera.center[:] = torch.tensor([W, H], dtype=dtype) * 0.5
# Re-enable gradient calculation for the camera translation
camera.translation.requires_grad = True
camera_opt_params = [camera.translation, body_model.global_orient]
elif camera_mode == 'fixed':
body_model.reset_params(body_pose=body_mean_pose, transl=init_t)
camera_opt_params = [body_model.transl, body_model.global_orient]
# If the distance between the 2D shoulders is smaller than a
# predefined threshold then try 2 fits, the initial one and a 180
# degree rotation
shoulder_dist = torch.dist(gt_joints[:, left_shoulder_idx],
gt_joints[:, right_shoulder_idx])
try_both_orient = shoulder_dist.item() < side_view_thsh
camera_optimizer, camera_create_graph = optim_factory.create_optimizer(
camera_opt_params, **kwargs)
# The closure passed to the optimizer
fit_camera = monitor.create_fitting_closure(
camera_optimizer,
body_model,
camera,
gt_joints,
camera_loss,
create_graph=camera_create_graph,
use_vposer=use_vposer,
vposer=vposer,
pose_embedding=pose_embedding,
scan_tensor=scan_tensor,
return_full_pose=False,
return_verts=False)
# Step 1: Optimize over the torso joints the camera translation
# Initialize the computational graph by feeding the initial translation
# of the camera and the initial pose of the body model.
camera_init_start = time.time()
cam_init_loss_val = monitor.run_fitting(camera_optimizer,
fit_camera,
camera_opt_params,
body_model,
use_vposer=use_vposer,
pose_embedding=pose_embedding,
vposer=vposer)
if interactive:
if use_cuda and torch.cuda.is_available():
torch.cuda.synchronize()
tqdm.write('Camera initialization done after {:.4f}'.format(
time.time() - camera_init_start))
tqdm.write('Camera initialization final loss {:.4f}'.format(
cam_init_loss_val))
# If the 2D detections/positions of the shoulder joints are too
# close the rotate the body by 180 degrees and also fit to that
# orientation
if try_both_orient:
body_orient = body_model.global_orient.detach().cpu().numpy()
flipped_orient = cv2.Rodrigues(body_orient)[0].dot(
cv2.Rodrigues(np.array([0., np.pi, 0]))[0])
flipped_orient = cv2.Rodrigues(flipped_orient)[0].ravel()
flipped_orient = torch.tensor(flipped_orient,
dtype=dtype,
device=device).unsqueeze(dim=0)
orientations = [body_orient, flipped_orient]
else:
orientations = [body_model.global_orient.detach().cpu().numpy()]
# store here the final error for both orientations,
# and pick the orientation resulting in the lowest error
results = []
body_transl = body_model.transl.clone().detach()
# Step 2: Optimize the full model
final_loss_val = 0
for or_idx, orient in enumerate(tqdm(orientations,
desc='Orientation')):
opt_start = time.time()
new_params = defaultdict(transl=body_transl,
global_orient=orient,
body_pose=body_mean_pose)
body_model.reset_params(**new_params)
if use_vposer:
with torch.no_grad():
pose_embedding.fill_(0)
for opt_idx, curr_weights in enumerate(
tqdm(opt_weights, desc='Stage')):
if opt_idx not in trans_opt_stages:
body_model.transl.requires_grad = False
else:
body_model.transl.requires_grad = True
body_params = list(body_model.parameters())
final_params = list(
filter(lambda x: x.requires_grad, body_params))
if use_vposer:
final_params.append(pose_embedding)
body_optimizer, body_create_graph = optim_factory.create_optimizer(
final_params, **kwargs)
body_optimizer.zero_grad()
curr_weights['bending_prior_weight'] = (
3.17 * curr_weights['body_pose_weight'])
if use_hands:
joint_weights[:, 25:76] = curr_weights['hand_weight']
if use_face:
joint_weights[:, 76:] = curr_weights['face_weight']
loss.reset_loss_weights(curr_weights)
closure = monitor.create_fitting_closure(
body_optimizer,
body_model,
camera=camera,
gt_joints=gt_joints,
joints_conf=joints_conf,
joint_weights=joint_weights,
loss=loss,
create_graph=body_create_graph,
use_vposer=use_vposer,
vposer=vposer,
pose_embedding=pose_embedding,
scan_tensor=scan_tensor,
scene_v=scene_v,
scene_vn=scene_vn,
scene_f=scene_f,
ftov=ftov,
return_verts=True,
return_full_pose=True)
if interactive:
if use_cuda and torch.cuda.is_available():
torch.cuda.synchronize()
stage_start = time.time()
final_loss_val = monitor.run_fitting(
body_optimizer,
closure,
final_params,
body_model,
pose_embedding=pose_embedding,
vposer=vposer,
use_vposer=use_vposer)
if interactive:
if use_cuda and torch.cuda.is_available():
torch.cuda.synchronize()
elapsed = time.time() - stage_start
if interactive:
tqdm.write(
'Stage {:03d} done after {:.4f} seconds'.format(
opt_idx, elapsed))
if interactive:
if use_cuda and torch.cuda.is_available():
torch.cuda.synchronize()
elapsed = time.time() - opt_start
tqdm.write(
'Body fitting Orientation {} done after {:.4f} seconds'.
format(or_idx, elapsed))
tqdm.write(
'Body final loss val = {:.5f}'.format(final_loss_val))
# Get the result of the fitting process
# Store in it the errors list in order to compare multiple
# orientations, if they exist
result = {
'camera_' + str(key): val.detach().cpu().numpy()
for key, val in camera.named_parameters()
}
result.update({
key: val.detach().cpu().numpy()
for key, val in body_model.named_parameters()
})
if use_vposer:
result['pose_embedding'] = pose_embedding.detach().cpu().numpy(
)
body_pose = vposer.decode(pose_embedding,
output_type='aa').view(
1, -1) if use_vposer else None
result['body_pose'] = body_pose.detach().cpu().numpy()
results.append({'loss': final_loss_val, 'result': result})
with open(result_fn, 'wb') as result_file:
if len(results) > 1:
min_idx = (0 if results[0]['loss'] < results[1]['loss'] else 1)
else:
min_idx = 0
pickle.dump(results[min_idx]['result'], result_file, protocol=2)
if save_meshes or visualize:
body_pose = vposer.decode(pose_embedding, output_type='aa').view(
1, -1) if use_vposer else None
model_type = kwargs.get('model_type', 'smpl')
append_wrists = model_type == 'smpl' and use_vposer
if append_wrists:
wrist_pose = torch.zeros([body_pose.shape[0], 6],
dtype=body_pose.dtype,
device=body_pose.device)
body_pose = torch.cat([body_pose, wrist_pose], dim=1)
model_output = body_model(return_verts=True, body_pose=body_pose)
vertices = model_output.vertices.detach().cpu().numpy().squeeze()
import trimesh
out_mesh = trimesh.Trimesh(vertices, body_model.faces, process=False)
out_mesh.export(mesh_fn)
if render_results:
import pyrender
# common
H, W = 1080, 1920
camera_center = np.array([951.30, 536.77])
camera_pose = np.eye(4)
camera_pose = np.array([1.0, -1.0, -1.0, 1.0]).reshape(-1,
1) * camera_pose
camera = pyrender.camera.IntrinsicsCamera(fx=1060.53,
fy=1060.38,
cx=camera_center[0],
cy=camera_center[1])
light = pyrender.DirectionalLight(color=np.ones(3), intensity=2.0)
material = pyrender.MetallicRoughnessMaterial(
metallicFactor=0.0,
alphaMode='OPAQUE',
baseColorFactor=(1.0, 1.0, 0.9, 1.0))
body_mesh = pyrender.Mesh.from_trimesh(out_mesh, material=material)
## rendering body
img = img.detach().cpu().numpy()
H, W, _ = img.shape
scene = pyrender.Scene(bg_color=[0.0, 0.0, 0.0, 0.0],
ambient_light=(0.3, 0.3, 0.3))
scene.add(camera, pose=camera_pose)
scene.add(light, pose=camera_pose)
# for node in light_nodes:
# scene.add_node(node)
scene.add(body_mesh, 'mesh')
r = pyrender.OffscreenRenderer(viewport_width=W,
viewport_height=H,
point_size=1.0)
color, _ = r.render(scene, flags=pyrender.RenderFlags.RGBA)
color = color.astype(np.float32) / 255.0
valid_mask = (color[:, :, -1] > 0)[:, :, np.newaxis]
input_img = img
output_img = (color[:, :, :-1] * valid_mask +
(1 - valid_mask) * input_img)
img = pil_img.fromarray((output_img * 255).astype(np.uint8))
img.save(out_img_fn)
##redering body+scene
body_mesh = pyrender.Mesh.from_trimesh(out_mesh, material=material)
static_scene = trimesh.load(osp.join(scene_dir, scene_name + '.ply'))
trans = np.linalg.inv(cam2world)
static_scene.apply_transform(trans)
static_scene_mesh = pyrender.Mesh.from_trimesh(static_scene)
scene = pyrender.Scene()
scene.add(camera, pose=camera_pose)
scene.add(light, pose=camera_pose)
scene.add(static_scene_mesh, 'mesh')
scene.add(body_mesh, 'mesh')
r = pyrender.OffscreenRenderer(viewport_width=W, viewport_height=H)
color, _ = r.render(scene)
color = color.astype(np.float32) / 255.0
img = pil_img.fromarray((color * 255).astype(np.uint8))
img.save(body_scene_rendering_fn)
def __init__(self, render):
# terms = 'f', 'frustum', 'background_image', 'overdraw', 'num_channels'
# dterms = 'vc', 'camera', 'bgcolor'
self.first_pass = True
self.render = render
self.scene = pyrender.Scene()
#self.human_mat = pyrender.MetallicRoughnessMaterial(baseColorFactor=[0.0, 0.0, 1.0 ,0.0])
self.human_mat = pyrender.MetallicRoughnessMaterial(
baseColorFactor=[0.05, 0.05, 0.8, 0.0]) #
self.human_mat_GT = pyrender.MetallicRoughnessMaterial(
baseColorFactor=[0.0, 0.3, 0.0, 0.0])
self.human_arm_mat = pyrender.MetallicRoughnessMaterial(
baseColorFactor=[0.1, 0.1, 0.8, 1.0])
self.human_mat_for_study = pyrender.MetallicRoughnessMaterial(
baseColorFactor=[0.3, 0.3, 0.3, 0.5])
self.human_bed_for_study = pyrender.MetallicRoughnessMaterial(
baseColorFactor=[0.7, 0.7, 0.2, 0.5])
self.human_mat_D = pyrender.MetallicRoughnessMaterial(
baseColorFactor=[0.1, 0.1, 0.1, 1.0], alphaMode="BLEND")
#if render == True:
mesh_color_mult = 0.25
self.mesh_parts_mat_list = [
pyrender.MetallicRoughnessMaterial(baseColorFactor=[
mesh_color_mult * 166. / 255., mesh_color_mult * 206. /
255., mesh_color_mult * 227. / 255., 0.0
]),
pyrender.MetallicRoughnessMaterial(baseColorFactor=[
mesh_color_mult * 31. / 255., mesh_color_mult * 120. /
255., mesh_color_mult * 180. / 255., 0.0
]),
pyrender.MetallicRoughnessMaterial(baseColorFactor=[
mesh_color_mult * 251. / 255., mesh_color_mult * 154. /
255., mesh_color_mult * 153. / 255., 0.0
]),
pyrender.MetallicRoughnessMaterial(baseColorFactor=[
mesh_color_mult * 227. / 255., mesh_color_mult * 26. /
255., mesh_color_mult * 28. / 255., 0.0
]),
pyrender.MetallicRoughnessMaterial(baseColorFactor=[
mesh_color_mult * 178. / 255., mesh_color_mult * 223. /
255., mesh_color_mult * 138. / 255., 0.0
]),
pyrender.MetallicRoughnessMaterial(baseColorFactor=[
mesh_color_mult * 51. / 255., mesh_color_mult * 160. /
255., mesh_color_mult * 44. / 255., 0.0
]),
pyrender.MetallicRoughnessMaterial(baseColorFactor=[
mesh_color_mult * 253. / 255., mesh_color_mult * 191. /
255., mesh_color_mult * 111. / 255., 0.0
]),
pyrender.MetallicRoughnessMaterial(baseColorFactor=[
mesh_color_mult * 255. / 255., mesh_color_mult * 127. /
255., mesh_color_mult * 0. / 255., 0.0
]),
pyrender.MetallicRoughnessMaterial(baseColorFactor=[
mesh_color_mult * 202. / 255., mesh_color_mult * 178. /
255., mesh_color_mult * 214. / 255., 0.0
]),
pyrender.MetallicRoughnessMaterial(baseColorFactor=[
mesh_color_mult * 106. / 255., mesh_color_mult * 61. /
255., mesh_color_mult * 154. / 255., 0.0
])
]
self.artag_mat = pyrender.MetallicRoughnessMaterial(
baseColorFactor=[0.3, 1.0, 0.3, 0.5])
self.artag_mat_other = pyrender.MetallicRoughnessMaterial(
baseColorFactor=[0.1, 0.1, 0.1, 0.0])
#self.artag_r = np.array([[-0.055, -0.055, 0.0], [-0.055, 0.055, 0.0], [0.055, -0.055, 0.0], [0.055, 0.055, 0.0]])
self.artag_r = np.array([
[0.0, 0.0, 0.075],
[0.0286 * 64 * 1.04 / 1.04, 0.0, 0.075],
[0.0, 0.01, 0.075],
[0.0286 * 64 * 1.04 / 1.04, 0.01, 0.075],
[0.0, 0.0, 0.075],
[0.0, 0.0286 * 27 / 1.06, 0.075],
[0.01, 0.0, 0.075],
[0.01, 0.0286 * 27 / 1.06, 0.075],
[0.0, 0.0286 * 27 / 1.06, 0.075],
[0.0286 * 64 * 1.04 / 1.04, 0.0286 * 27 / 1.06, 0.075],
[0.0, 0.0286 * 27 / 1.06 + 0.01, 0.075],
[0.0286 * 64 * 1.04 / 1.04, 0.0286 * 27 / 1.06 + 0.01, 0.075],
[0.0286 * 64 * 1.04 / 1.04, 0.0, 0.075],
[0.0286 * 64 * 1.04 / 1.04, 0.0286 * 27 / 1.06, 0.075],
[0.0286 * 64 * 1.04 / 1.04 - 0.01, 0.0, 0.075],
[0.0286 * 64 * 1.04 / 1.04 - 0.01, 0.0286 * 27 / 1.06, 0.075],
])
#self.artag_f = np.array([[0, 1, 3], [3, 1, 0], [0, 2, 3], [3, 2, 0], [1, 3, 2]])
self.artag_f = np.array([[0, 1, 2], [0, 2, 1], [1, 2, 3], [1, 3, 2],
[4, 5, 6], [4, 6, 5], [5, 6, 7], [5, 7, 6],
[8, 9, 10], [8, 10, 9], [9, 10, 11],
[9, 11, 10], [12, 13, 14], [12, 14, 13],
[13, 14, 15], [13, 15, 14]])
#self.artag_facecolors_root = np.array([[0.0, 1.0, 0.0],[0.0, 1.0, 0.0],[0.0, 1.0, 0.0],[0.0, 1.0, 0.0],[0.0, 1.0, 0.0]])
self.artag_facecolors_root = np.array([
[0.3, 0.3, 0.0],
[0.3, 0.3, 0.0],
[0.3, 0.3, 0.0],
[0.3, 0.3, 0.0],
[0.3, 0.3, 0.0],
[0.3, 0.3, 0.0],
[0.3, 0.3, 0.0],
[0.3, 0.3, 0.0],
[0.3, 0.3, 0.0],
[0.3, 0.3, 0.0],
[0.3, 0.3, 0.0],
[0.3, 0.3, 0.0],
[0.3, 0.3, 0.0],
[0.3, 0.3, 0.0],
[0.3, 0.3, 0.0],
[0.3, 0.3, 0.0],
])
#self.artag_facecolors = np.array([[0.0, 0.0, 0.0],[0.0, 0.0, 0.0],[0.0, 0.0, 0.0],[0.0, 0.0, 0.0],[0.0, 0.0, 0.0],])
self.artag_facecolors = np.copy(self.artag_facecolors_root)
self.pic_num = 0
def render_image(img,
verts,
cam,
faces=None,
angle=None,
axis=None,
resolution=224,
output_fn=None):
if faces is None:
faces = get_smpl_faces()
mesh = trimesh.Trimesh(vertices=verts, faces=faces)
Rx = trimesh.transformations.rotation_matrix(math.radians(180), [1, 0, 0])
mesh.apply_transform(Rx)
if angle and axis:
R = trimesh.transformations.rotation_matrix(math.radians(angle), axis)
mesh.apply_transform(R)
if output_fn:
mesh.export(output_fn)
camera_translation = np.array(
[-cam[1], cam[2], 2 * 5000. / (img.shape[0] * cam[0] + 1e-9)])
np.save(output_fn.replace('.obj', '.npy'), camera_translation)
# Save the rotated mesh
# R = trimesh.transformations.rotation_matrix(math.radians(270), [0,1,0])
# rotated_mesh = mesh.copy()
# rotated_mesh.apply_transform(R)
# rotated_mesh.export(output_fn.replace('.obj', '_rot.obj'))
scene = pyrender.Scene(bg_color=[0.0, 0.0, 0.0, 0.0],
ambient_light=(0.3, 0.3, 0.3))
material = pyrender.MetallicRoughnessMaterial(metallicFactor=0.0,
alphaMode='OPAQUE',
baseColorFactor=(1.0, 1.0,
0.9, 1.0))
mesh = pyrender.Mesh.from_trimesh(mesh, material=material)
scene.add(mesh, 'mesh')
camera_pose = np.eye(4)
camera = WeakPerspectiveCamera(scale=cam[0],
translation=cam[1:],
zfar=1000.)
scene.add(camera, pose=camera_pose)
light = pyrender.PointLight(color=[1.0, 1.0, 1.0], intensity=1)
light_pose = np.eye(4)
light_pose[:3, 3] = [0, -1, 1]
scene.add(light, pose=light_pose)
light_pose[:3, 3] = [0, 1, 1]
scene.add(light, pose=light_pose)
light_pose[:3, 3] = [1, 1, 2]
scene.add(light, pose=light_pose)
r = pyrender.OffscreenRenderer(viewport_width=resolution,
viewport_height=resolution,
point_size=1.0)
color, _ = r.render(scene, flags=pyrender.RenderFlags.RGBA)
# color = color[:, ::-1, :]
valid_mask = (color[:, :, -1] > 0)[:, :, np.newaxis]
output_img = color[:, :, :-1] * valid_mask + (1 - valid_mask) * img
image = output_img.astype(np.uint8)
text = f's: {cam[0]:.2f}, tx: {cam[1]:.2f}, ty: {cam[2]:.2f}'
cv2.putText(image, text, (5, 10), 0, 0.4, color=(0, 255, 0))
return image
def __call__(self,
verts,
img=np.zeros((224, 224, 3)),
cam=np.array([1, 0, 0]),
angle=None,
axis=None,
mesh_filename=None,
color_type=None,
color=[1.0, 1.0, 0.9]):
mesh = trimesh.Trimesh(vertices=verts, faces=self.faces, process=False)
Rx = trimesh.transformations.rotation_matrix(math.radians(180),
[1, 0, 0])
mesh.apply_transform(Rx)
if mesh_filename is not None:
mesh.export(mesh_filename)
if angle and axis:
R = trimesh.transformations.rotation_matrix(
math.radians(angle), axis)
mesh.apply_transform(R)
if len(cam) == 4:
sx, sy, tx, ty = cam
elif len(cam) == 3:
sx, tx, ty = cam
sy = sx
camera = WeakPerspectiveCamera(scale=[sx, sy],
translation=[tx, ty],
zfar=1000.)
if color_type != None:
color = self.colors_dict[color_type]
material = pyrender.MetallicRoughnessMaterial(
metallicFactor=0.0,
alphaMode='OPAQUE',
baseColorFactor=(color[0], color[1], color[2], 1.0))
mesh = pyrender.Mesh.from_trimesh(mesh, material=material)
mesh_node = self.scene.add(mesh, 'mesh')
camera_pose = np.eye(4)
cam_node = self.scene.add(camera, pose=camera_pose)
if self.wireframe:
render_flags = RenderFlags.RGBA | RenderFlags.ALL_WIREFRAME
else:
render_flags = RenderFlags.RGBA
rgb, _ = self.renderer.render(self.scene, flags=render_flags)
valid_mask = (rgb[:, :, -1] > 0)[:, :, np.newaxis]
image_list = [img] if type(img) is not list else img
return_img = []
for item in image_list:
output_img = rgb[:, :, :-1] * valid_mask + (1 - valid_mask) * item
image = output_img.astype(np.uint8)
return_img.append(image)
if type(img) is not list:
return_img = return_img[0]
self.scene.remove_node(mesh_node)
self.scene.remove_node(cam_node)
return return_img
def render_multiview(self, vertices, K, R, T, imglist, trackId=0, return_depth=False, return_color=False,
bg_color=[0.0, 0.0, 0.0, 0.0], camera=None):
# List to store rendered scenes
output_images, output_colors, output_depths = [], [], []
# Need to flip x-axis
rot = trimesh.transformations.rotation_matrix(
np.radians(180), [1, 0, 0])
nViews = len(imglist)
for nv in range(nViews):
img = imglist[nv]
self.renderer.viewport_height = img.shape[0]
self.renderer.viewport_width = img.shape[1]
# Create a scene for each image and render all meshes
scene = pyrender.Scene(bg_color=bg_color,
ambient_light=(0.3, 0.3, 0.3))
camera_pose = np.eye(4)
# for every person in the scene
if isinstance(vertices, dict):
for trackId, data in vertices.items():
vert = data['vertices'].copy()
faces = data['faces']
col = data.get('col', trackId)
vert = vert @ R[nv].T + T[nv]
mesh = trimesh.Trimesh(vert, faces)
mesh.apply_transform(rot)
trans = [0, 0, 0]
material = pyrender.MetallicRoughnessMaterial(
metallicFactor=0.0,
alphaMode='OPAQUE',
baseColorFactor=colors[col % len(colors)])
mesh = pyrender.Mesh.from_trimesh(
mesh,
material=material)
scene.add(mesh, 'mesh')
else:
verts = vertices @ R[nv].T + T[nv]
mesh = trimesh.Trimesh(verts, self.faces)
mesh.apply_transform(rot)
trans = [0, 0, 0]
material = pyrender.MetallicRoughnessMaterial(
metallicFactor=0.0,
alphaMode='OPAQUE',
baseColorFactor=colors[trackId % len(colors)])
mesh = pyrender.Mesh.from_trimesh(
mesh,
material=material)
scene.add(mesh, 'mesh')
if camera is not None:
light = pyrender.PointLight(color=[1.0, 1.0, 1.0], intensity=70)
light_pose = np.eye(4)
light_pose[:3, 3] = [0, 0, 4.5]
scene.add(light, pose=light_pose)
light_pose[:3, 3] = [0, 1, 4]
scene.add(light, pose=light_pose)
light_pose[:3, 3] = [0, -1, 4]
scene.add(light, pose=light_pose)
else:
trans = [0, 0, 0]
# Use 3 directional lights
# Create light source
light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=1)
light_pose = np.eye(4)
light_pose[:3, 3] = np.array([0, -1, 1]) + trans
scene.add(light, pose=light_pose)
light_pose[:3, 3] = np.array([0, 1, 1]) + trans
scene.add(light, pose=light_pose)
light_pose[:3, 3] = np.array([1, 1, 2]) + trans
scene.add(light, pose=light_pose)
if camera is None:
if K is None:
camera_center = np.array([img.shape[1] / 2., img.shape[0] / 2.])
camera = pyrender.camera.IntrinsicsCamera(fx=self.focal_length, fy=self.focal_length, cx=camera_center[0], cy=camera_center[1])
else:
camera = pyrender.camera.IntrinsicsCamera(fx=K[nv][0, 0], fy=K[nv][1, 1], cx=K[nv][0, 2], cy=K[nv][1, 2])
scene.add(camera, pose=camera_pose)
# Alpha channel was not working previously need to check again
# Until this is fixed use hack with depth image to get the opacity
color, rend_depth = self.renderer.render(scene, flags=flags)
# color = color[::-1,::-1]
# rend_depth = rend_depth[::-1,::-1]
output_depths.append(rend_depth)
color = color.astype(np.uint8)
valid_mask = (rend_depth > 0)[:, :, None]
if color.shape[2] == 3: # 在服务器上透明通道失败
color = np.dstack((color, (valid_mask*255).astype(np.uint8)))
output_colors.append(color)
output_img = (color[:, :, :3] * valid_mask +
(1 - valid_mask) * img)
output_img = output_img.astype(np.uint8)
output_images.append(output_img)
if return_depth:
return output_images, output_depths
elif return_color:
return output_colors
else:
return output_images
def show_scene(vertices,
faces,
camera_pose,
image,
K,
joints=[],
color=(0.8, 0.3, 0.3, 1.0)):
mats = []
for c in color:
material = pyrender.MetallicRoughnessMaterial(metallicFactor=0.2,
alphaMode='OPAQUE',
baseColorFactor=c)
mats.append(material)
scene = pyrender.Scene(ambient_light=(0.5, 0.5, 0.5))
for i, v in enumerate(vertices):
mesh = trimesh.Trimesh(v, faces)
rot = trimesh.transformations.rotation_matrix(np.radians(180),
[1, 0, 0])
mesh.apply_transform(rot)
mesh = pyrender.Mesh.from_trimesh(mesh, material=mats[i])
scene.add(mesh, 'mesh')
camera = pyrender.IntrinsicsCamera(fx=K[0, 0],
fy=K[1, 1],
cx=K[0, 2],
cy=K[1, 2])
scene.add(camera, pose=camera_pose)
cam = trimesh.creation.axis()
mesh = pyrender.Mesh.from_trimesh(cam, smooth=False)
scene.add(mesh, pose=camera_pose)
scene.add(mesh, pose=np.linalg.inv(camera_pose))
scene.add(mesh, pose=np.eye(4))
light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=1)
light_pose = np.eye(4)
light_pose[:3, 3] = np.array([0, -1, 1])
scene.add(light, pose=light_pose)
light_pose[:3, 3] = np.array([0, 1, 1])
scene.add(light, pose=light_pose)
light_pose[:3, 3] = np.array([1, 1, 2])
scene.add(light, pose=light_pose)
for i, j_list in enumerate(joints):
for j in j_list:
mesh = trimesh.creation.uv_sphere(0.01, count=[32, 32])
mesh = pyrender.Mesh.from_trimesh(mesh, material=mats[i])
pos = np.eye(4)
pos[:3, 3] = j[:-1]
scene.add(mesh, 'mesh', pose=pos)
pyrender.Viewer(scene)
