Spaces:
Running
on
L40S
Running
on
L40S
| 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 | |