import os
import cv2
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import matplotlib as mpl
import os
os.environ['PYOPENGL_PLATFORM'] = 'egl'
import pyrender
import trimesh
from config.config import cfg
def vis_keypoints_with_skeleton(img, kps, kps_lines, kp_thresh=0.4, alpha=1):
# Convert from plt 0-1 RGBA colors to 0-255 BGR colors for opencv.
cmap = plt.get_cmap('rainbow')
colors = [cmap(i) for i in np.linspace(0, 1, len(kps_lines) + 2)]
colors = [(c[2] * 255, c[1] * 255, c[0] * 255) for c in colors]
# Perform the drawing on a copy of the image, to allow for blending.
kp_mask = np.copy(img)
# Draw the keypoints.
for l in range(len(kps_lines)):
i1 = kps_lines[l][0]
i2 = kps_lines[l][1]
p1 = kps[0, i1].astype(np.int32), kps[1, i1].astype(np.int32)
p2 = kps[0, i2].astype(np.int32), kps[1, i2].astype(np.int32)
if kps[2, i1] > kp_thresh and kps[2, i2] > kp_thresh:
cv2.line(kp_mask,
p1,
p2,
color=colors[l],
thickness=2,
lineType=cv2.LINE_AA)
if kps[2, i1] > kp_thresh:
cv2.circle(kp_mask,
p1,
radius=3,
color=colors[l],
thickness=-1,
lineType=cv2.LINE_AA)
if kps[2, i2] > kp_thresh:
cv2.circle(kp_mask,
p2,
radius=3,
color=colors[l],
thickness=-1,
lineType=cv2.LINE_AA)
# Blend the keypoints.
return cv2.addWeighted(img, 1.0 - alpha, kp_mask, alpha, 0)
def vis_keypoints(img, kps, alpha=1, radius=3, color=None):
# Convert from plt 0-1 RGBA colors to 0-255 BGR colors for opencv.
cmap = plt.get_cmap('rainbow')
if color is None:
colors = [cmap(i) for i in np.linspace(0, 1, len(kps) + 2)]
colors = [(c[2] * 255, c[1] * 255, c[0] * 255) for c in colors]
# Perform the drawing on a copy of the image, to allow for blending.
kp_mask = np.copy(img)
# Draw the keypoints.
for i in range(len(kps)):
p = kps[i][0].astype(np.int32), kps[i][1].astype(np.int32)
if color is None:
cv2.circle(kp_mask,
p,
radius=radius,
color=colors[i],
thickness=-1,
lineType=cv2.LINE_AA)
else:
cv2.circle(kp_mask,
p,
radius=radius,
color=color,
thickness=-1,
lineType=cv2.LINE_AA)
# Blend the keypoints.
return cv2.addWeighted(img, 1.0 - alpha, kp_mask, alpha, 0)
def vis_mesh(img, mesh_vertex, alpha=0.5):
# Convert from plt 0-1 RGBA colors to 0-255 BGR colors for opencv.
cmap = plt.get_cmap('rainbow')
colors = [cmap(i) for i in np.linspace(0, 1, len(mesh_vertex))]
colors = [(c[2] * 255, c[1] * 255, c[0] * 255) for c in colors]
# Perform the drawing on a copy of the image, to allow for blending.
mask = np.copy(img)
# Draw the mesh
for i in range(len(mesh_vertex)):
p = mesh_vertex[i][0].astype(np.int32), mesh_vertex[i][1].astype(
np.int32)
cv2.circle(mask,
p,
radius=1,
color=colors[i],
thickness=-1,
lineType=cv2.LINE_AA)
# Blend the keypoints.
return cv2.addWeighted(img, 1.0 - alpha, mask, alpha, 0)
def vis_3d_skeleton(kpt_3d, kpt_3d_vis, kps_lines, filename=None):
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
# Convert from plt 0-1 RGBA colors to 0-255 BGR colors for opencv.
cmap = plt.get_cmap('rainbow')
colors = [cmap(i) for i in np.linspace(0, 1, len(kps_lines) + 2)]
colors = [np.array((c[2], c[1], c[0])) for c in colors]
for l in range(len(kps_lines)):
i1 = kps_lines[l][0]
i2 = kps_lines[l][1]
x = np.array([kpt_3d[i1, 0], kpt_3d[i2, 0]])
y = np.array([kpt_3d[i1, 1], kpt_3d[i2, 1]])
z = np.array([kpt_3d[i1, 2], kpt_3d[i2, 2]])
if kpt_3d_vis[i1, 0] > 0 and kpt_3d_vis[i2, 0] > 0:
ax.plot(x, z, -y, c=colors[l], linewidth=2)
if kpt_3d_vis[i1, 0] > 0:
ax.scatter(kpt_3d[i1, 0],
kpt_3d[i1, 2],
-kpt_3d[i1, 1],
c=colors[l],
marker='o')
if kpt_3d_vis[i2, 0] > 0:
ax.scatter(kpt_3d[i2, 0],
kpt_3d[i2, 2],
-kpt_3d[i2, 1],
c=colors[l],
marker='o')
x_r = np.array([0, cfg.input_shape[1]], dtype=np.float32)
y_r = np.array([0, cfg.input_shape[0]], dtype=np.float32)
z_r = np.array([0, 1], dtype=np.float32)
if filename is None:
ax.set_title('3D vis')
else:
ax.set_title(filename)
ax.set_xlabel('X Label')
ax.set_ylabel('Z Label')
ax.set_zlabel('Y Label')
ax.legend()
plt.show()
cv2.waitKey(0)
def save_obj(v, f, file_name='output.obj'):
obj_file = open(file_name, 'w')
for i in range(len(v)):
obj_file.write('v ' + str(v[i][0]) + ' ' + str(v[i][1]) + ' ' +
str(v[i][2]) + '\n')
for i in range(len(f)):
obj_file.write('f ' + str(f[i][0] + 1) + '/' + str(f[i][0] + 1) + ' ' +
str(f[i][1] + 1) + '/' + str(f[i][1] + 1) + ' ' +
str(f[i][2] + 1) + '/' + str(f[i][2] + 1) + '\n')
obj_file.close()
def perspective_projection(vertices, cam_param):
# vertices: [N, 3]
# cam_param: [3]
fx, fy = cam_param['focal']
cx, cy = cam_param['princpt']
vertices[:, 0] = vertices[:, 0] * fx / vertices[:, 2] + cx
vertices[:, 1] = vertices[:, 1] * fy / vertices[:, 2] + cy
return vertices
def render_mesh(img, mesh, face, cam_param, mesh_as_vertices=False):
if mesh_as_vertices:
# to run on cluster where headless pyrender is not supported for A100/V100
vertices_2d = perspective_projection(mesh, cam_param)
img = vis_keypoints(img,
vertices_2d,
alpha=0.8,
radius=2,
color=(0, 0, 255))
else:
# 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