Spaces:
Runtime error
Runtime error
| ''' | |
| NIMBLE: A Non-rigid Hand Model with Bones and Muscles[SIGGRAPH-22] | |
| https://reyuwei.github.io/proj/nimble | |
| ''' | |
| import torch | |
| import pytorch3d | |
| import numpy as np | |
| from pathlib import Path | |
| from pytorch3d.structures.meshes import Meshes | |
| import pytorch3d.ops | |
| ROOT_JOINT_IDX = 0 # wrist | |
| DOF2_BONES = [1, 2, 4, 5, 8, 9, 12, 13, 16, 17] | |
| DOF1_BONES = [3, 6, 7, 10, 11, 14, 15, 18, 19] | |
| JOINT_PARENT_ID_DICT = { | |
| 0: -1, | |
| 1: 0, | |
| 2: 1, | |
| 3: 2, | |
| 4: 3, | |
| 5: 0, | |
| 6: 5, | |
| 7: 6, | |
| 8: 7, | |
| 9: 8, | |
| 10: 0, | |
| 11: 10, | |
| 12: 11, | |
| 13: 12, | |
| 14: 13, | |
| 15: 0, | |
| 16: 15, | |
| 17: 16, | |
| 18: 17, | |
| 19: 18, | |
| 20: 0, | |
| 21: 20, | |
| 22: 21, | |
| 23: 22, | |
| 24: 23 | |
| } | |
| JOINT_ID_NAME_DICT = { | |
| 0: "carpal", | |
| 1: "met1", | |
| 2: "pro1", | |
| 3: "dis1", | |
| 4: "dis1_end", | |
| 5: "met2", | |
| 6: "pro2", | |
| 7: "int2", | |
| 8: "dis2", | |
| 9: "dis2_end", | |
| 10: "met3", | |
| 11: "pro3", | |
| 12: "int3", | |
| 13: "dis3", | |
| 14: "dis3_end", | |
| 15: "met4", | |
| 16: "pro4", | |
| 17: "int4", | |
| 18: "dis4", | |
| 19: "dis4_end", | |
| 20: "met5", | |
| 21: "pro5", | |
| 22: "int5", | |
| 23: "dis5", | |
| 24: "dis5_end" | |
| } | |
| BONE_TO_JOINT_NAME = { | |
| 0: "carpal", | |
| 1: "met1", | |
| 2: "pro1", | |
| 3: "dis1", | |
| 4: "met2", | |
| 5: "pro2", | |
| 6: "int2", | |
| 7: "dis2", | |
| 8: "met3", | |
| 9: "pro3", | |
| 10: "int3", | |
| 11: "dis3", | |
| 12: "met4", | |
| 13: "pro4", | |
| 14: "int4", | |
| 15: "dis4", | |
| 16: "met5", | |
| 17: "pro5", | |
| 18: "int5", | |
| 19: "dis5", | |
| } | |
| STATIC_BONE_NUM = 20 | |
| STATIC_JOINT_NUM = 25 | |
| JOINT_ID_BONE_DICT = {} | |
| JOINT_ID_BONE = np.zeros(STATIC_BONE_NUM) | |
| BONE_ID_JOINT_DICT = {} | |
| for key in JOINT_ID_NAME_DICT: | |
| value = JOINT_ID_NAME_DICT[key] | |
| for key_b in BONE_TO_JOINT_NAME: | |
| if BONE_TO_JOINT_NAME[key_b] == value: | |
| JOINT_ID_BONE_DICT[key] = key_b | |
| BONE_ID_JOINT_DICT[key_b] = key | |
| JOINT_ID_BONE[key_b] = key | |
| def dis_to_weight(dismat, thres_corres, node_sigma): | |
| dismat[dismat==0] = 1e5 | |
| dismat[dismat>thres_corres] = 1e5 | |
| node_weight = torch.exp(-dismat / node_sigma) | |
| norm = torch.norm(node_weight, dim=1) | |
| norm_node_weight = node_weight / (norm + 1e-6) | |
| norm_node_weight[norm==0] = 0 | |
| return norm_node_weight | |
| def batch_to_tensor_device(batch, device): | |
| def to_tensor(arr): | |
| if isinstance(arr, int): | |
| return arr | |
| if isinstance(arr, torch.Tensor): | |
| return arr.to(device) | |
| if arr.dtype == np.int64: | |
| arr = torch.from_numpy(arr) | |
| else: | |
| arr = torch.from_numpy(arr).float() | |
| return arr | |
| for key in batch: | |
| if isinstance(batch[key], np.ndarray): | |
| batch[key] = to_tensor(batch[key]).to(device) | |
| elif isinstance(batch[key], list): | |
| for i in range(len(batch[key])): | |
| if isinstance(batch[key][i], list): | |
| for j in range(len(batch[key][i])): | |
| if isinstance(batch[key][i][j], np.ndarray): | |
| batch[key][i][j] = to_tensor(batch[key][i][j]).to(device) | |
| else: | |
| batch[key][i] = to_tensor(batch[key][i]).to(device) | |
| elif isinstance(batch[key], dict): | |
| batch[key] = batch_to_tensor_device(batch[key], device) | |
| elif isinstance(batch[key], torch.Tensor): | |
| batch[key] = batch[key].to(device) | |
| return batch | |
| def quat2aa(quats): | |
| """ | |
| Convert wxyz quaternions to angle-axis representation | |
| :param quats: | |
| :return: | |
| """ | |
| _cos = quats[..., 0] | |
| xyz = quats[..., 1:] | |
| _sin = xyz.norm(dim=-1) | |
| norm = _sin.clone() | |
| norm[norm < 1e-7] = 1 | |
| axis = xyz / norm.unsqueeze(-1) | |
| angle = torch.atan2(_sin, _cos) * 2 | |
| return axis * angle.unsqueeze(-1) | |
| def quat2mat(quat): | |
| """Convert quaternion coefficients to rotation matrix. | |
| Args: | |
| quat: size = [batch_size, 4] 4 <===>(w, x, y, z) | |
| Returns: | |
| Rotation matrix corresponding to the quaternion -- size = [batch_size, 3, 3] | |
| """ | |
| norm_quat = quat | |
| norm_quat = norm_quat / norm_quat.norm(p=2, dim=1, keepdim=True) | |
| w, x, y, z = norm_quat[:, 0], norm_quat[:, 1], norm_quat[:, | |
| 2], norm_quat[:, | |
| 3] | |
| batch_size = quat.size(0) | |
| w2, x2, y2, z2 = w.pow(2), x.pow(2), y.pow(2), z.pow(2) | |
| wx, wy, wz = w * x, w * y, w * z | |
| xy, xz, yz = x * y, x * z, y * z | |
| rotMat = torch.stack([ | |
| w2 + x2 - y2 - z2, 2 * xy - 2 * wz, 2 * wy + 2 * xz, 2 * wz + 2 * xy, | |
| w2 - x2 + y2 - z2, 2 * yz - 2 * wx, 2 * xz - 2 * wy, 2 * wx + 2 * yz, | |
| w2 - x2 - y2 + z2 | |
| ], | |
| dim=1).view(batch_size, 3, 3) | |
| return rotMat | |
| def batch_aa2quat(axisang): | |
| # w, x, y, z | |
| axisang_norm = torch.norm(axisang + 1e-8, p=2, dim=1) | |
| angle = torch.unsqueeze(axisang_norm, -1) | |
| axisang_normalized = torch.div(axisang, angle) | |
| angle = angle * 0.5 | |
| v_cos = torch.cos(angle) | |
| v_sin = torch.sin(angle) | |
| quat = torch.cat([v_cos, v_sin * axisang_normalized], dim=1) | |
| return quat | |
| def batch_rodrigues(axisang): | |
| #axisang N x 3 | |
| axisang_norm = torch.norm(axisang + 1e-8, p=2, dim=1) | |
| angle = torch.unsqueeze(axisang_norm, -1) | |
| axisang_normalized = torch.div(axisang, angle) | |
| angle = angle * 0.5 | |
| v_cos = torch.cos(angle) | |
| v_sin = torch.sin(angle) | |
| quat = torch.cat([v_cos, v_sin * axisang_normalized], dim=1) | |
| rot_mat = quat2mat(quat) | |
| rot_mat = rot_mat.view(rot_mat.shape[0], 9) | |
| return rot_mat | |
| def th_posemap_axisang_2output(pose_vectors): | |
| rot_nb = int(pose_vectors.shape[1] / 3) | |
| rot_mats = [] | |
| for joint_idx in range(rot_nb - 1): | |
| joint_idx_val = joint_idx + 1 | |
| axis_ang = pose_vectors[:, joint_idx_val * 3:(joint_idx_val + 1) * 3] | |
| rot_mat = batch_rodrigues(axis_ang) | |
| rot_mats.append(rot_mat) | |
| # rot_mats = torch.stack(rot_mats, 1).view(-1, 15 *9) | |
| rot_mats = torch.cat(rot_mats, 1) | |
| pose_maps = subtract_flat_id(rot_mats) | |
| return pose_maps, rot_mats | |
| def subtract_flat_id(rot_mats): | |
| # Subtracts identity as a flattened tensor | |
| rot_nb = int(rot_mats.shape[1] / 9) | |
| id_flat = torch.eye( | |
| 3, dtype=rot_mats.dtype, device=rot_mats.device).view(1, 9).repeat( | |
| rot_mats.shape[0], rot_nb) | |
| # id_flat.requires_grad = False | |
| results = rot_mats - id_flat | |
| return results | |
| def th_with_zeros(tensor): | |
| batch_size = tensor.shape[0] | |
| padding = tensor.new([0.0, 0.0, 0.0, 1.0]) | |
| padding.requires_grad = False | |
| concat_list = [tensor, padding.view(1, 1, 4).repeat(batch_size, 1, 1)] | |
| cat_res = torch.cat(concat_list, 1) | |
| return cat_res | |
| def th_scalemat_scale(th_scale_bone): | |
| batch_size = th_scale_bone.shape[0] | |
| th_scale_bone_mat = torch.eye(4).repeat([batch_size, th_scale_bone.shape[1], 1, 1]) | |
| th_scale_bone_mat = th_scale_bone_mat.type_as(th_scale_bone).to(th_scale_bone.device) | |
| if len(th_scale_bone.shape) == 3: | |
| for s in range(th_scale_bone.shape[1]): | |
| th_scale_bone_mat[:, s, 0, 0] = th_scale_bone[:, s, 0] | |
| th_scale_bone_mat[:, s, 1, 1] = th_scale_bone[:, s, 1] | |
| th_scale_bone_mat[:, s, 2, 2] = th_scale_bone[:, s, 2] | |
| else: | |
| for s in range(th_scale_bone.shape[1]): | |
| th_scale_bone_mat[:, s, 0, 0] = th_scale_bone[:, s] | |
| th_scale_bone_mat[:, s, 1, 1] = th_scale_bone[:, s] | |
| th_scale_bone_mat[:, s, 2, 2] = th_scale_bone[:, s] | |
| return th_scale_bone_mat | |
| def th_pack(tensor): | |
| batch_size = tensor.shape[0] | |
| padding = tensor.new_zeros((batch_size, 4, 3)) | |
| padding.requires_grad = False | |
| pack_list = [padding, tensor] | |
| pack_res = torch.cat(pack_list, 2) | |
| return pack_res | |
| def vertices2landmarks( | |
| vertices, | |
| faces, | |
| lmk_faces_idx, | |
| lmk_bary_coords | |
| ): | |
| ''' | |
| Calculates landmarks by barycentric interpolation | |
| Parameters | |
| ---------- | |
| vertices: torch.tensor BxVx3, dtype = torch.float32 | |
| The tensor of input vertices | |
| faces: torch.tensor Fx3, dtype = torch.long | |
| The faces of the mesh | |
| lmk_faces_idx: torch.tensor L, dtype = torch.long | |
| The tensor with the indices of the faces used to calculate the | |
| landmarks. | |
| lmk_bary_coords: torch.tensor Lx3, dtype = torch.float32 | |
| The tensor of barycentric coordinates that are used to interpolate | |
| the landmarks | |
| Returns | |
| ------- | |
| landmarks: torch.tensor BxLx3, dtype = torch.float32 | |
| The coordinates of the landmarks for each mesh in the batch | |
| Modified from https://github.com/vchoutas/smplx | |
| ''' | |
| # Extract the indices of the vertices for each face | |
| # BxLx3 | |
| batch_size, num_verts = vertices.shape[:2] | |
| device = vertices.device | |
| # lmk_faces = torch.index_select(faces, 0, lmk_faces_idx.view(-1)).view( | |
| # batch_size, -1, 3) | |
| lmk_faces = torch.index_select(faces, 0, lmk_faces_idx.view(-1)).view( | |
| 1, -1, 3) | |
| lmk_faces = lmk_faces.repeat([batch_size,1,1]) | |
| lmk_faces += torch.arange( | |
| batch_size, dtype=torch.long, device=device).view(-1, 1, 1) * num_verts | |
| lmk_vertices = vertices.reshape(-1, 3)[lmk_faces].view( | |
| batch_size, -1, 3, 3) | |
| landmarks = torch.einsum('blfi,lf->bli', [lmk_vertices, lmk_bary_coords]) | |
| return landmarks | |
| def save_textured_nimble(fname, skin_v, tex_img): | |
| ### batch_size = 1 | |
| import cv2 | |
| textured_pkl = "assets/NIMBLE_TEX_FUV.pkl" | |
| fname = Path(fname) | |
| obj_name_skin = fname.parent / (fname.stem + "_skin.obj") | |
| mtl_name = obj_name_skin.with_suffix(".mtl") | |
| # texture image | |
| tex_name_diffuse = fname.parent / (fname.stem + "_diffuse.png") | |
| tex_img = np.uint8(tex_img * 255) | |
| cv2.imwrite(str(tex_name_diffuse), tex_img[:,:, :3]) | |
| cv2.imwrite(str(fname.parent / (fname.stem + "_normal.png")), tex_img[:,:,3:6]) | |
| cv2.imwrite(str(fname.parent / (fname.stem + "_spec.png")), tex_img[:,:,6:]) | |
| # mtl | |
| mtl_str = "newmtl material_0\nKa 0.200000 0.200000 0.200000\nKd 0.800000 0.800000 0.800000\nKs 1.000000 1.000000 1.000000\nTr 1.000000\nillum 2\nNs 0.000000\nmap_Kd " | |
| mtl_str = mtl_str + tex_name_diffuse.name | |
| with open(mtl_name, "w") as f: | |
| f.writelines(mtl_str) | |
| # obj | |
| f_uv = np.load(textured_pkl, allow_pickle=True) | |
| with open(obj_name_skin, "w") as f: | |
| f.write("mtllib {:s}\n".format(mtl_name.name)) | |
| for v in skin_v: | |
| f.writelines("v {:.5f} {:.5f} {:.5f}\n".format(v[0], v[1], v[2])) | |
| f.writelines(f_uv) | |
| print("save to", fname) | |
| def smooth_mesh(mesh_p3d): | |
| mesh_p3d_smooth = pytorch3d.ops.mesh_filtering.taubin_smoothing(mesh_p3d, num_iter=3) | |
| target_mv = mesh_p3d_smooth.verts_padded() | |
| nan_mv = torch.isnan(target_mv) | |
| target_mv[nan_mv] = mesh_p3d.verts_padded()[nan_mv] | |
| mesh_p3d_smooth_fixnan = Meshes(target_mv, mesh_p3d.faces_padded()) | |
| return mesh_p3d_smooth_fixnan | |