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 render_glcam(K,
Rt,
model_in, # model name or trimesh
scale=1.0,
std_size=(1000, 1000),
flat_shading=False):
# Mesh creation
if isinstance(model_in, str) is True:
mesh = trimesh.load(model_in)
else:
mesh = model_in.copy()
pr_mesh = pyrender.Mesh.from_trimesh(mesh)
# Scene creation
scene = pyrender.Scene()
# Adding objects to the scene
face_node = scene.add(pr_mesh)
# Caculate fx fy cx cy from K
fx, fy = K[0][0] * scale, K[1][1] * scale
cx, cy = K[0][2] * scale, K[1][2] * scale
# Camera Creation
cam = pyrender.IntrinsicsCamera(fx, fy, cx, cy,
znear=0.1, zfar=100000)
cam_pose = np.eye(4)
cam_pose[:3, :3] = Rt[:3, :3].T
cam_pose[:3, 3] = -Rt[:3, :3].T.dot(Rt[:, 3])
scene.add(cam, pose=cam_pose)
# Set up the light
light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=10.0)
scene.add(light, pose=cam_pose)
# Rendering offscreen from that camera
r = pyrender.OffscreenRenderer(viewport_width=std_size[1],
viewport_height=std_size[0],
point_size=1.0)
if flat_shading is True:
color, depth = r.render(scene, flags=pyrender.constants.RenderFlags.FLAT)
else:
color, depth = r.render(scene)
# rgb to bgr for cv2
color = color[:, :, [2, 1, 0]]
return depth, color
def __call__(self, vertices, camera_translation, image):
material = pyrender.MetallicRoughnessMaterial(
metallicFactor=0.0,
alphaMode='OPAQUE',
smooth=False,
wireframe=True,
roughnessFactor=1.0,
emissiveFactor=(0.1, 0.1, 0.1),
baseColorFactor=(1.0, 1.0, 0.9, 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)
mesh = pyrender.Mesh.from_trimesh(mesh,
material=material,
smooth=False)
scene = pyrender.Scene(ambient_light=(0.30, 0.30, 0.30))
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=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)
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 face_rendering(mesh, camera_pose, light_poses, show=True):
"""
Render face RGBD images with input camera pose and lighting
:param mesh: Trimesh object
:param camera_pose: Twc, np.array 4x4
:param light_poses: list of light poses, Twc, list[np.array 4x4]
:param show: whether show rendered image
:return:
"""
mesh = pyrender.Mesh.from_trimesh(mesh)
scene = pyrender.Scene()
scene.add(mesh)
# Set up the camera -- z-axis away from the scene, x-axis right, y-axis up
camera = pyrender.PerspectiveCamera(yfov=np.pi / 10.0)
scene.add(camera, pose=camera_pose)
# Set up the light
for light_pose in light_poses:
light = pyrender.DirectionalLight(color=np.ones(3), intensity=10.0)
light_pose = rotation_matrix(angle=0.0, direction=[0.0, 1.0, 0.0])
scene.add(light, pose=light_pose)
# Render the scene
r = pyrender.OffscreenRenderer(960, 1280)
color, depth = r.render(scene)
# depth[depth < 1e-5] = 0.75
# Show the images
if show:
img_list = [{'img': color, 'title': 'RGB'},
{'img': depth, 'title': 'Depth'}]
show_multiple_img(img_list, num_cols=2)
# print(depth[480, 640])
r.delete()
# Compute camera pose Twc
Twc = camera_pose
T = np.array([
[1.0, 0.0, 0.0, 0.0],
[0.0, -1.0, 0.0, 0.0],
[0.0, 0.0, -1.0, 0.0],
[0.0, 0.0, 0.0, 1.0],
])
Twc = np.dot(T, np.dot(Twc, T))
return color, depth, K_from_PerspectiveCamera(camera, 1280, 960), Twc
def simple_face_rendering(obj_file_path, show=True):
mesh = load_mesh_from_obj(obj_file_path)
mesh = pyrender.Mesh.from_trimesh(mesh)
scene = pyrender.Scene()
scene.add(mesh)
# Set up the camera -- z-axis away from the scene, x-axis right, y-axis up
camera = pyrender.PerspectiveCamera(yfov=np.pi / 10.0)
camera_pose = np.array([
[1.0, 0.0, 0.0, 0.0],
[0.0, 1.0, 0.0, 0.0],
[0.0, 0.0, 1.0, 1.0],#/300],
[0.0, 0.0, 0.0, 1.0],
])
# camera_pose = rotation_matrix(angle=np.pi / 4.0, direction=[0.0, 1.0, 0.0])
# camera_pose[0, 3] = camera_pose[2, 3] = np.sqrt(2) / 2
scene.add(camera, pose=camera_pose)
# Set up the light -- a single spot light in the same spot as the camera
light = pyrender.DirectionalLight(color=np.ones(3), intensity=10.0)
light_pose = rotation_matrix(angle=0.0, direction=[0.0, 1.0, 0.0])
scene.add(light, pose=light_pose)
# Render the scene
r = pyrender.OffscreenRenderer(960, 1280)
color, depth = r.render(scene)
# depth[depth < 1e-5] = 0.75
# Show the images
if show:
img_list = [{'img': color, 'title': 'RGB'},
{'img': depth, 'title': 'Depth'}]
show_multiple_img(img_list, num_cols=2)
# print(depth[480, 640])
r.delete()
# Compute camera pose Twc
Twc = camera_pose
T = np.array([
[1.0, 0.0, 0.0, 0.0],
[0.0, -1.0, 0.0, 0.0],
[0.0, 0.0, -1.0, 0.0],
[0.0, 0.0, 0.0, 1.0],
])
Twc = np.dot(T, np.dot(Twc, T))
return color, depth, K_from_PerspectiveCamera(camera, 1280, 960), Twc
def _instantiate_pyrender_scene(self):
self._scene = pyrender.Scene()
self._renderer = pyrender.OffscreenRenderer(self.image_size[0],
self.image_size[1])
light = pyrender.DirectionalLight(color=np.ones(3), intensity=1.0)
cam = pyrender.IntrinsicsCamera(
fx=self.focal_length[0],
fy=self.focal_length[0],
cx=self.image_size[0] / 2,
cy=self.image_size[1] / 2,
zfar=100000000000000 # `Infinite` clipping
)
self._light_obj = self._scene.add(light)
self._cam_obj = self._scene.add(cam)
def __init__(self, img_size, bg_color=None):
if bg_color is None:
bg_color = np.array([0.1, 0.1, 0.1, 1.])
self.scene = pyrender.Scene(bg_color=bg_color)
self.focal_len = 5.
camera = pyrender.PerspectiveCamera(
yfov=np.tan(1 / self.focal_len) * 2, aspectRatio=1.0)
camera_pose = np.eye(4, dtype=np.float32)
self.scene.add(camera, pose=camera_pose)
light = pyrender.DirectionalLight(
color=np.ones(3),
intensity=10.0,
)
self.scene.add(light, pose=camera_pose)
if not hasattr(img_size, '__iter__'):
img_size = [img_size, img_size]
self.r = pyrender.OffscreenRenderer(*img_size)
def render_face_orthographic(mesh, background=None):
"""
mesh location should be normalized
:param mesh:
:param background:
:return:
"""
mesh.visual.face_colors = np.array([0.05, 0.1, 0.2, 1])
mesh = pyrender.Mesh.from_trimesh(mesh, smooth=False)
# mesh = pyrender.Mesh.from_trimesh(mesh)
scene.add(mesh, pose=np.eye(4))
camera_pose = np.eye(4)
# camera_pose[0, 3] = 1
# camera_pose[1, 3] = 1
# camera_pose[2, 3] = -10
# camera_pose[0, 0] = 1
# camera_pose[1, 1] = -1
# camera_pose[2, 2] = -1
#
# camera = pyrender.OrthographicCamera(xmag=1, ymag=1, zfar=100)
camera_pose[0, 3] = 1
camera_pose[1, 3] = 1
camera_pose[2, 3] = 10
camera_pose[0, 0] = 1
camera_pose[1, 1] = 1
camera_pose[2, 2] = 1
camera = pyrender.OrthographicCamera(xmag=1, ymag=1, zfar=100)
scene.add(camera, pose=camera_pose)
light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=5.0)
scene.add(light, pose=camera_pose)
color, depth = r.render(scene)
scene.clear()
# print(color.shape)
color = np.array(color)
color = color[::-1]
if background is not None:
new_color = np.array(background)
new_color[color != 255] = color[color != 255]
color = new_color
return color
def renderLight(posmap, init_image=None, is_render=True):
tex = np.ones((256, 256, 3)) / 2
mesh = UVmap2Mesh(posmap, tex, is_extra_triangle=False)
vertices = mesh['vertices']
triangles = mesh['triangles']
colors = mesh['colors'] / np.max(mesh['colors'])
file = 'tmp/light/test.obj'
write_obj_with_colors(file, vertices, triangles, colors)
obj = trimesh.load(file)
# obj.visual.vertex_colors = np.random.uniform(size=obj.vertices.shape)
obj.visual.face_colors = np.array([0.05, 0.1, 0.2])
mesh = pyrender.Mesh.from_trimesh(obj, smooth=False)
scene.add(mesh, pose=np.eye(4))
camera_pose = np.eye(4)
camera_pose[0, 3] = 128
camera_pose[1, 3] = 128
camera_pose[2, 3] = 300
camera = pyrender.OrthographicCamera(xmag=128, ymag=128, zfar=1000)
scene.add(camera, pose=camera_pose)
light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=8.0)
scene.add(light, pose=camera_pose)
color, depth = r.render(scene)
if is_render:
plt.imshow(color)
plt.show()
if init_image is not None:
sum_mask = np.mean(color, axis=-1)
fuse_img = color.copy()
fuse_img[sum_mask > 128] = init_image[sum_mask > 128]
if is_render:
plt.imshow(fuse_img)
plt.show()
scene.clear()
return fuse_img
scene.clear()
return color
def __init__(self,
skeleton_name,
mode='key',
camera_pose=None,
camera_intrin='1280_720_color',
suppress_warnings=False):
super().__init__(skeleton_name)
intrin = makeIntrinsics(camera_intrin)
self.mode = mode
self.suppress_warnings = suppress_warnings
ml = MeshLoader()
ml.load()
name_list = ml.getNames()
self.meshes = ml.getMeshes()
if camera_pose is not None:
c_pose = camera_pose
else:
c_pose = [.087, -1.425, .4, 0, 1.551, -.025]
self.scene = pyrender.Scene(bg_color=[0.0, 0.0, 0.0]) # Make scene
camera = cameraFromIntrinsics(intrin)
cam_pose = makePose(*c_pose)
self.camera_node = self.scene.add(camera, pose=cam_pose)
dl = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=10.0)
self.scene.add(
dl, parent_node=self.camera_node) # Add light at camera pose
# Add in joints
self.joint_nodes = []
for mesh, name in zip(self.meshes, name_list):
self.joint_nodes.append(pyrender.Node(name=name, mesh=mesh))
for node in self.joint_nodes:
self.scene.add_node(node)
self._updateKeypoints()
self.rend = pyrender.OffscreenRenderer(intrin.width, intrin.height)
self.node_color_map = {}
self.setMode(mode)
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 setup_standard_scene(camera):
"""
Creates an empty scene with some standard lighting and the given camera.
Parameters
----------
camera:
pyrender camera that should be used for rendering
Returns
-------
an empty scene with lighting and camera set up.
"""
directional_light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=2.0)
scene = pyrender.Scene()
scene.add(camera)
scene.add(directional_light)
return scene
def render_mesh(model):
dist = 2
angle = 20
height = -0.2
scene = pyrender.Scene(ambient_light=[.1, .1, .3], bg_color=(0, 0, 0))
scene.add(pyrender.Mesh.from_trimesh(model, smooth=False))
light = pyrender.DirectionalLight(color=[1, 1, 1], intensity=2e3)
scene.add(light, pose=np.eye(4))
c = np.cos(angle * np.pi / 180)
s = np.sin(angle * np.pi / 180)
camera_pose = np.array([[c, 0, s, dist * s], [0, 1, 0, height],
[-1 * s, 0, c, dist * c], [0, 0, 0, 1]])
camera = pyrender.PerspectiveCamera(yfov=np.pi / 3.0, znear=0.5, zfar=5)
scene.add(camera, pose=camera_pose)
renderer = pyrender.OffscreenRenderer(512, 512)
color, _ = renderer.render(scene)
return color[:, :, ::-1]
def export_3d_image(self, filename="3d_render.png", tolerance=0.005):
"""Creates a 3D rendered image (png) of the reactor
:param filename: output filename of the image created
:type filename: [ParamType](, optional)
:param tolerance: the mesh tolerance
:type tolerance: float
:return: filename of the created image
:rtype: str
"""
scene = pyrender.Scene(ambient_light=np.array([0.1, 0.1, 0.1, 1.0]))
for entry in self.shapes_and_components:
if entry.render_mesh is None:
scene.add(entry._create_render_mesh(tolerance))
# sets the field of view (fov) and the aspect ratio of the image
camera = pyrender.camera.PerspectiveCamera(yfov=math.radians(90.0),
aspectRatio=2.0)
# sets the position of the camera using a matrix
c = 2**-0.5
camera_pose = np.array([[1, 0, 0, 0], [0, c, -c, -500], [0, c, c, 500],
[0, 0, 0, 1]])
scene.add(camera, pose=camera_pose)
# adds some basic lighting to the scene
light = pyrender.DirectionalLight(color=np.ones(3), intensity=1.0)
scene.add(light, pose=camera_pose)
# Render the scene
renderer = pyrender.OffscreenRenderer(1000, 500)
colours, depth = renderer.render(scene)
image = Image.fromarray(colours, "RGB")
Path(filename).parent.mkdir(parents=True, exist_ok=True)
image.save(filename, "PNG")
print("\n saved 3d image to ", filename)
return filename
def save_image_face_heatmap(facedata, vec, errors, id, name="face_heat", path=""):
plt.clf()
plt.close()
colors_heat = []
# Map errors to RGB colors
min_error = 0
max_error = 6
for i, er in enumerate(errors[id]):
c_er = abs(er)
if c_er > max_error:
c_er = max_error
c = rgb(min_error, max_error, c_er)
colors_heat.append(c)
# Generate colored mesh
predict_trimeshh = facedata.vec2meshTrimesh2(vec, col=colors_heat)
trimeshh = pyrender.Mesh.from_trimesh(predict_trimeshh, smooth=False)
# Create scene for rendering, etc.
scene = pyrender.Scene(ambient_light=[.1, .1, .3], bg_color=[255, 255, 255])
camera = pyrender.PerspectiveCamera(yfov=np.pi / 3.0)
light = pyrender.DirectionalLight(color=[1, 1, 1], intensity=2e3)
scene.add(trimeshh, pose=np.eye(4))
scene.add(light, pose=np.eye(4))
# non chiedermi mai come ho trovato questi valori bellissimi (T/N: just use these values)
camera_pose = np.array([
[0.94063, 0.01737, -0.41513, -88.15790],
[-0.06728, 0.98841, -0.16663, -35.36127],
[0.33266, 0.15078, 1.14014, 241.71166],
[0.00000, 0.00000, 0.00000, 1.00000]
])
scene.add(camera, pose=camera_pose)
r = pyrender.OffscreenRenderer(512, 512)
color, _ = r.render(scene)
plt.figure(figsize=(8, 8))
plt.imshow(color)
name_image = path+name+str(".png")
# plt.colorbar() # Could be a nice addition, however unneeded
plt.savefig(name_image)
plt.clf()
def __init__(self):
"""Construct a Scene."""
super().__init__()
self._bullet_nodes = {}
self._seg_node_map = {}
self.bg_color = (0.7, 0.7, 0.8)
self.ambient_light = (0.2, 0.2, 0.2)
self._camera_node = pyr.Node(camera=pyr.PerspectiveCamera(
np.deg2rad(60.0)),
translation=(0.0, -2.0, 3.0),
rotation=(-0.472, 0.0, 0.0, 0.882))
self.add_node(self._camera_node)
self._light_node = pyr.Node(light=pyr.DirectionalLight(color=(0.8, 0.8,
0.8),
intensity=5.0),
translation=(-0.8, -0.2, 2.0),
rotation=(-0.438, 0.342, -0.511, 0.655))
self.add_node(self._light_node)
def test2():
model_path = 'data/models/basic/cow.obj'
# load the cow model
tm = trimesh.load(model_path)
tm.visual.vertex_colors = np.random.uniform(
size=tm.visual.vertex_colors.shape)
tm.visual.face_colors = np.random.uniform(size=tm.visual.face_colors.shape)
mesh = pyrender.Mesh.from_trimesh(tm, smooth=False)
light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=2.0)
cam = pyrender.PerspectiveCamera(yfov=np.pi / 3.0, aspectRatio=1.414)
nm = pyrender.Node(mesh=mesh, matrix=np.eye(4))
nl = pyrender.Node(light=light, matrix=np.eye(4))
nc = pyrender.Node(camera=cam, matrix=np.eye(4))
scene = pyrender.Scene(ambient_light=[1.0, 1.0, 1.0], bg_color=gray)
scene.add_node(nm)
scene.add_node(nl, parent_node=nm)
scene.add_node(nc, parent_node=nm)
pyrender.Viewer(scene, use_raymond_lighting=True)
def _add_lighting(scene, light_type, random_range=(1, 4)):
'''Takes scene and adds random amout of lighting.
random_range: Range to pick number of lights from.'''
n = random.randrange( # Number of lights
random_range[0], random_range[1])
for _ in range(n): # Add directional lights
d = None
if 'directional_lights' == light_type:
d = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=2)
elif 'point_lights' == light_type:
d = pyrender.PointLight(color=[1.0, 1.0, 1.0], intensity=2)
else:
raise Exception(
'Light type not recognized, should be \"direction_lights\" or \"point_lights\", not {}'
.format(light_type))
_add_model(scene, d)
return scene
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 img_renderer(img, mesh, light_est):
# compose scene
scene = pyrender.Scene(ambient_light=[.1, .1, .3], bg_color=[0, 0, 0])
camera = pyrender.PerspectiveCamera(yfov=np.pi / 3.0)
light = pyrender.DirectionalLight(color=[1, 1, 1], intensity=2e3)
scene.add(mesh, pose=np.eye(4))
scene.add(light_est, pose=np.eye(4))
# c = 2**-0.5
# scene.add(camera, pose=[[ 1, 0, 0, 0],
# [ 0, c, -c, -2],
# [ 0, c, c, 2],
# [ 0, 0, 0, 1]])
# render scene
r = pyrender.OffscreenRenderer(512, 512)
color, _ = r.render(scene)
plt.figure(figsize=(8, 8)), plt.imshow(color)
def render_example():
# generate mesh
sphere = trimesh.creation.icosphere(subdivisions=4, radius=0.8)
sphere.vertices += 1e-2 * np.random.randn(*sphere.vertices.shape)
mesh = pyrender.Mesh.from_trimesh(sphere, smooth=False)
# compose scene
scene = pyrender.Scene(ambient_light=[.1, .1, .3], bg_color=[0, 0, 0])
camera = pyrender.PerspectiveCamera(yfov=np.pi / 3.0)
light = pyrender.DirectionalLight(color=[1, 1, 1], intensity=2e3)
scene.add(mesh, pose=np.eye(4))
scene.add(light, pose=np.eye(4))
c = 2**-0.5
scene.add(camera,
pose=[[1, 0, 0, 0], [0, c, -c, -2], [0, c, c, 2], [0, 0, 0, 1]])
# render scene
r = pyrender.OffscreenRenderer(512, 512)
color, _ = r.render(scene)
plt.figure(figsize=(8, 8)), plt.imshow(color)
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 __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
for meantheta in np.linspace(mintheta, maxtheta, anglesteps):
for meanphi in np.linspace(minphi, maxphi, anglesteps):
theta = meantheta + np.random.uniform(
-(maxtheta - mintheta) / anglesteps / 2,
(maxtheta - mintheta) / anglesteps / 2)
phi = meanphi + np.random.uniform(
-(maxphi - minphi) / anglesteps / 2,
(maxphi - minphi) / anglesteps / 2)
dist = meandist + np.random.uniform(
-(maxdist - mindist) / diststeps / 2,
(maxdist - mindist) / diststeps / 2)
scene = pyrender.Scene(ambient_light=[.1, .1, .3],
bg_color=[1, 1, 1])
camera = pyrender.PerspectiveCamera(yfov=np.pi / 3.0)
intensity = 10**np.random.uniform(3, 3.5)
light = pyrender.DirectionalLight(color=[1, 1, 1],
intensity=intensity)
scene.add(mesh, pose=np.eye(4))
x = np.random.uniform(0, 5)
y = np.random.uniform(0, 5)
z = np.random.uniform(0, 5)
scene.add(light,
pose=[[1, 0, 0, x], [0, 1, 0, y], [0, 0, 1, z],
[0, 0, 0, 1]])
ct = np.cos(theta)
st = np.sin(theta)
cp = np.cos(phi)
sp = np.sin(phi)
hor_rotation = np.array([[ct, 0, st, 0], [0, 1, 0, 0],
[-st, 0, ct, 0], [0, 0, 0, 1]])
vert_rotation = np.array([[1, 0, 0, 0], [0, cp, -sp, 0],
[0, sp, cp, 0], [0, 0, 0, 1]])
def render_orthcam(
model_in, # model name or trimesh
xy_mag,
rend_size,
flat_shading=False,
zfar=10000,
znear=0.05):
# Mesh creation
if isinstance(model_in, str) is True:
mesh = trimesh.load(model_in, process=False)
else:
mesh = model_in.copy()
pr_mesh = pyrender.Mesh.from_trimesh(mesh)
# Scene creation
scene = pyrender.Scene()
# Adding objects to the scene
face_node = scene.add(pr_mesh)
# Camera Creation
if type(xy_mag) == float:
cam = pyrender.OrthographicCamera(xmag=xy_mag,
ymag=xy_mag,
znear=znear,
zfar=zfar)
elif type(xy_mag) == tuple:
cam = pyrender.OrthographicCamera(xmag=xy_mag[0],
ymag=xy_mag[1],
znear=znear,
zfar=zfar)
else:
print("Error: xy_mag should be float or tuple")
return False
scene.add(cam, pose=np.eye(4))
# Set up the light
light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=10.0)
scene.add(light, pose=np.eye(4))
# Rendering offscreen from that camera
r = pyrender.OffscreenRenderer(viewport_width=rend_size[1],
viewport_height=rend_size[0],
point_size=1.0)
if flat_shading is True:
color, depth = r.render(scene,
flags=pyrender.constants.RenderFlags.FLAT)
else:
color, depth = r.render(scene)
# rgb to bgr for cv2
color = color[:, :, [2, 1, 0]]
# fix pyrender BUG of depth rendering, pyrender version: 0.1.43
depth[depth != 0] = (zfar + znear - (
(2.0 * znear * zfar) / depth[depth != 0])) / (zfar - znear)
depth[depth != 0] = ((depth[depth != 0] + (zfar + znear) /
(zfar - znear)) * (zfar - znear)) / 2.0
return depth, color
dtype=np.float32).reshape(2)
princpt = np.array(cam_param['princpt'][cam_idx],
dtype=np.float32).reshape(2)
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_width,
viewport_height=img_height,
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)
depth = depth[:, :, None]
valid_mask = (depth > 0)
if prev_depth is None:
render_mask = valid_mask
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 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)
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
def render_smpl_on_image(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 * alpha + \
valid_mask * image / 255 * (1-alpha) + (1 - valid_mask) * image / 255
cv2.imwrite(filename, (255 * output_img).astype(np.int16))
