iti/HandMesh
0
1from __future__ import absolute_import2from __future__ import division3from __future__ import print_function4from __future__ import unicode_literals5 6import cv27import numpy as np8 9import matplotlib10 11matplotlib.use('Agg')12import matplotlib.pyplot as plt13from mpl_toolkits.mplot3d import Axes3D14import matplotlib.tri as mtri15 16color_hand_joints = [[1.0, 0.0, 0.0],17 [0.0, 0.4, 0.0], [0.0, 0.6, 0.0], [0.0, 0.8, 0.0], [0.0, 1.0, 0.0], # thumb18 [0.0, 0.0, 0.6], [0.0, 0.0, 1.0], [0.2, 0.2, 1.0], [0.4, 0.4, 1.0], # index19 [0.0, 0.4, 0.4], [0.0, 0.6, 0.6], [0.0, 0.8, 0.8], [0.0, 1.0, 1.0], # middle20 [0.4, 0.4, 0.0], [0.6, 0.6, 0.0], [0.8, 0.8, 0.0], [1.0, 1.0, 0.0], # ring21 [0.4, 0.0, 0.4], [0.6, 0.0, 0.6], [0.8, 0.0, 0.8], [1.0, 0.0, 1.0]] # little22 23camera_shape = [[[-0.05, 0.05, 0.05, -0.05, -0.05], [-0.05, -0.05, 0.05, 0.05, -0.05], [0, 0, 0, 0, 0]],24 [[0.05, 0], [0.05, 0], [0, -0.1]],25 [[0.05, 0], [-0.05, 0], [0, -0.1]],26 [[-0.05, 0], [-0.05, 0], [0, -0.1]],27 [[-0.05, 0], [0.05, 0], [0, -0.1]]28 ]29 30camera_color = (0, 0, 200/255)31 32 33def fig2data(fig):34 """35 @brief Convert a Matplotlib figure to a 4D numpy array with RGBA channels and return it36 @param fig a matplotlib figure37 @return a numpy 3D array of RGBA values38 """39 # draw the renderer40 fig.canvas.draw()41 42 # Get the RGBA buffer from the figure43 w, h = fig.canvas.get_width_height()44 buf = np.fromstring(fig.canvas.tostring_argb(), dtype=np.uint8)45 buf.shape = (w, h, 4)46 47 # canvas.tostring_argb give pixmap in ARGB mode. Roll the ALPHA channel to have it in RGBA mode48 buf = np.roll(buf, 3, axis=2)49 return buf50 51 52def draw_silhouette(image, mask=None, poly=None):53 """54 :param image: H x W x 355 :param mask: H x W56 :param poly: 1 x N x 2 (np.array)57 :return:58 """59 img_mask = image.copy()60 if mask is not None:61 mask = np.concatenate([np.zeros(list(mask.shape) + [2]), mask[:, :, None]], 2).astype(np.uint8) * 25562 img_mask = cv2.addWeighted(img_mask, 1, mask, 0.5, 0)63 if poly is not None:64 cv2.polylines(img_mask, poly, isClosed=True, thickness=2, color=(0, 0, 255))65 66 return img_mask67 68 69def draw_mesh(image, cam_param, mesh_xyz, face):70 """71 :param image: H x W x 372 :param cam_param: 1 x 3 x 373 :param mesh_xyz: 778 x 374 :param face: 1538 x 3 x 275 :return:76 """77 vertex2uv = np.matmul(cam_param, mesh_xyz.T).T78 vertex2uv = (vertex2uv / vertex2uv[:, 2:3])[:, :2].astype(np.int)79 80 fig = plt.figure()81 fig.set_size_inches(float(image.shape[0]) / fig.dpi, float(image.shape[1]) / fig.dpi, forward=True)82 plt.imshow(image)83 plt.axis('off')84 if face is None:85 plt.plot(vertex2uv[:, 0], vertex2uv[:, 1], 'o', color='green', markersize=1)86 else:87 plt.triplot(vertex2uv[:, 0], vertex2uv[:, 1], face, lw=0.5, color='orange')88 89 plt.subplots_adjust(left=0., right=1., top=1., bottom=0, wspace=0, hspace=0)90 91 ret = fig2data(fig)92 plt.close(fig)93 94 return ret95 96def draw_2d_skeleton(image, pose_uv):97 """98 :param image: H x W x 399 :param pose_uv: 21 x 2100 wrist,101 thumb_mcp, thumb_pip, thumb_dip, thumb_tip102 index_mcp, index_pip, index_dip, index_tip,103 middle_mcp, middle_pip, middle_dip, middle_tip,104 ring_mcp, ring_pip, ring_dip, ring_tip,105 little_mcp, little_pip, little_dip, little_tip106 :return:107 """108 assert pose_uv.shape[0] == 21109 skeleton_overlay = image.copy()110 marker_sz = 6111 line_wd = 3112 root_ind = 0113 114 for joint_ind in range(pose_uv.shape[0]):115 joint = pose_uv[joint_ind, 0].astype('int32'), pose_uv[joint_ind, 1].astype('int32')116 cv2.circle(117 skeleton_overlay, joint,118 radius=marker_sz, color=color_hand_joints[joint_ind] * np.array(255), thickness=-1,119 lineType=cv2.CV_AA if cv2.__version__.startswith('2') else cv2.LINE_AA)120 if joint_ind == 0:121 continue122 elif joint_ind % 4 == 1:123 root_joint = pose_uv[root_ind, 0].astype('int32'), pose_uv[root_ind, 1].astype('int32')124 cv2.line(125 skeleton_overlay, root_joint, joint,126 color=color_hand_joints[joint_ind] * np.array(255), thickness=int(line_wd),127 lineType=cv2.CV_AA if cv2.__version__.startswith('2') else cv2.LINE_AA)128 else:129 joint_2 = pose_uv[joint_ind - 1, 0].astype('int32'), pose_uv[joint_ind - 1, 1].astype('int32')130 cv2.line(131 skeleton_overlay, joint_2, joint,132 color=color_hand_joints[joint_ind] * np.array(255), thickness=int(line_wd),133 lineType=cv2.CV_AA if cv2.__version__.startswith('2') else cv2.LINE_AA)134 135 136 return skeleton_overlay137 138 139def draw_3d_skeleton(pose_cam_xyz, image_size):140 """141 :param pose_cam_xyz: 21 x 3142 :param image_size: H, W143 :return:144 """145 assert pose_cam_xyz.shape[0] == 21146 fig = plt.figure()147 fig.set_size_inches(float(image_size[0]) / fig.dpi, float(image_size[1]) / fig.dpi, forward=True)148 149 ax = plt.subplot(111, projection='3d')150 marker_sz = 10151 line_wd = 2152 153 for i, shape in enumerate(camera_shape):154 ax.plot(shape[0], shape[1], shape[2], color=camera_color, linestyle=(':', '-')[i==0])155 156 for joint_ind in range(pose_cam_xyz.shape[0]):157 ax.plot(pose_cam_xyz[joint_ind:joint_ind + 1, 0], pose_cam_xyz[joint_ind:joint_ind + 1, 1],158 pose_cam_xyz[joint_ind:joint_ind + 1, 2], '.', c=color_hand_joints[joint_ind], markersize=marker_sz)159 if joint_ind == 0:160 continue161 elif joint_ind % 4 == 1:162 ax.plot(pose_cam_xyz[[0, joint_ind], 0], pose_cam_xyz[[0, joint_ind], 1], pose_cam_xyz[[0, joint_ind], 2],163 color=color_hand_joints[joint_ind], linewidth=line_wd)164 else:165 ax.plot(pose_cam_xyz[[joint_ind - 1, joint_ind], 0], pose_cam_xyz[[joint_ind - 1, joint_ind], 1],166 pose_cam_xyz[[joint_ind - 1, joint_ind], 2], color=color_hand_joints[joint_ind],167 linewidth=line_wd)168 169 ax.axis('auto')170 x_lim = [-0.1, 0.1, 0.02]171 y_lim = [-0.1, 0.12, 0.02]172 z_lim = [0.0, 0.8, 0.1]173 x_ticks = np.arange(x_lim[0], x_lim[1], step=x_lim[2])174 y_ticks = np.arange(y_lim[0], y_lim[1], step=y_lim[2])175 z_ticks = np.arange(z_lim[0], z_lim[1], step=z_lim[2])176 plt.xticks(x_ticks, [x_lim[0], '', '', '', '', 0, '', '', '', x_lim[1]], fontsize=14)177 plt.yticks(y_ticks, [y_lim[0], '', '', '', '', 0, '', '', '', -y_lim[0], ''], fontsize=14)178 ax.set_zticks(z_ticks)179 z_ticks = [''] * (z_ticks.shape[0])180 z_ticks[4] = 0.4181 ax.set_zticklabels(z_ticks, fontsize=14)182 ax.view_init(elev=140, azim=80)183 plt.subplots_adjust(left=-0.06, right=0.98, top=0.93, bottom=-0.07, wspace=0, hspace=0)184 185 ret = fig2data(fig)186 plt.close(fig)187 return ret188 189def draw_3d_mesh(mesh_xyz, image_size, face):190 """191 :param mesh_xyz: 778 x 3192 :param image_size: H, W193 :param face: 1538 x 3194 :return:195 """196 fig = plt.figure()197 fig.set_size_inches(float(image_size[0]) / fig.dpi, float(image_size[1]) / fig.dpi, forward=True)198 199 ax = plt.subplot(111, projection='3d')200 201 for i, shape in enumerate(camera_shape):202 ax.plot(shape[0], shape[1], shape[2], color=camera_color, linestyle=(':', '-')[i==0])203 204 triang = mtri.Triangulation(mesh_xyz[:, 0], mesh_xyz[:, 1], triangles=face)205 ax.plot_trisurf(triang, mesh_xyz[:, 2], color=(145/255, 181/255, 255/255))206 207 ax.axis('auto')208 x_lim = [-0.1, 0.1, 0.02]209 y_lim = [-0.1, 0.12, 0.02]210 z_lim = [0.0, 0.8, 0.1]211 x_ticks = np.arange(x_lim[0], x_lim[1], step=x_lim[2])212 y_ticks = np.arange(y_lim[0], y_lim[1], step=y_lim[2])213 z_ticks = np.arange(z_lim[0], z_lim[1], step=z_lim[2])214 plt.xticks(x_ticks, [x_lim[0], '', '', '', '', 0, '', '', '', x_lim[1]], fontsize=14)215 plt.yticks(y_ticks, [y_lim[0], '', '', '', '', 0, '', '', '', -y_lim[0], ''], fontsize=14)216 ax.set_zticks(z_ticks)217 z_ticks = ['']*(z_ticks.shape[0])218 z_ticks[4] = 0.4219 ax.set_zticklabels(z_ticks, fontsize=14)220 ax.view_init(elev=140, azim=80)221 plt.subplots_adjust(left=-0.06, right=1, top=0.95, bottom=-0.06, wspace=0, hspace=0)222 plt.subplots_adjust(left=-0.06, right=0.98, top=0.93, bottom=-0.07, wspace=0, hspace=0)223 224 225 ret = fig2data(fig)226 plt.close(fig)227 return ret228 229def save_a_image_with_mesh_joints(image, mask, poly, cam_param, mesh_xyz, face, pose_uv, pose_xyz, file_name, padding=0, ret=False):230 """231 :param mesh_plot:232 :param image: H x W x 3 (np.array)233 :param mask: H x W (np.array)234 :param poly: 1 x N x 2 (np.array)235 :param cam_params: 3 x 3 (np.array)236 :param mesh_xyz: 778 x 3 (np.array)237 :param face: 1538 x 3 (np.array)238 :param pose_uv: 21 x 2 (np.array)239 :param pose_xyz: 21 x 3 (np.array)240 :param file_name:241 :param padding:242 :return:243 """244 if poly is not None:245 img_mask = draw_silhouette(image, mask, poly)246 else:247 img_mask = image.copy()248 rend_img_overlay = draw_mesh(image, cam_param, mesh_xyz, face)249 skeleton_overlay = draw_2d_skeleton(image, pose_uv)250 skeleton_3d = draw_3d_skeleton(pose_xyz, image.shape[:2])251 mesh_3d = draw_3d_mesh(mesh_xyz, image.shape[:2], face)252 253 img_list = [img_mask, skeleton_overlay, rend_img_overlay, mesh_3d, skeleton_3d]254 image_height = image.shape[0]255 image_width = image.shape[1]256 num_column = len(img_list)257 258 grid_image = np.zeros(((image_height + padding), num_column * (image_width + padding), 3), dtype=np.uint8)259 260 width_begin = 0261 width_end = image_width262 for show_img in img_list:263 grid_image[:, width_begin:width_end, :] = show_img[..., :3]264 width_begin += (image_width + padding)265 width_end = width_begin + image_width266 if ret:267 return grid_image268 269 cv2.imwrite(file_name, grid_image)270