Spaces:
Runtime error
Runtime error
| from numpy import ndarray | |
| from typing import Optional, Tuple | |
| import numpy as np | |
| import scipy | |
| def assert_ndarray(arr, name: str="arr", shape: Optional[Tuple[int, ...]]=None, dtype=None): | |
| if not isinstance(arr, np.ndarray): | |
| raise ValueError(f"{name} must be a numpy.ndarray or None, got {type(arr)}") | |
| if shape is not None: | |
| # shape may contain None as wildcard | |
| if len(shape) != arr.ndim: | |
| raise ValueError(f"{name}: expected shape length {len(shape)} but array ndim is {arr.ndim}") | |
| for i, (exp, actual) in enumerate(zip(shape, arr.shape)): | |
| if exp > 0 and exp != actual: | |
| raise ValueError(f"{name} shape mismatch at axis {i}: expected {exp}, got {actual}") | |
| if dtype is not None: | |
| if not np.issubdtype(arr.dtype, dtype): | |
| raise ValueError(f"{name} dtype must be {dtype}, got {arr.dtype}") | |
| def assert_list(arr, name: str="arr", dtype=None): | |
| if not isinstance(arr, list): | |
| raise ValueError(f"found type {type(arr)}, expect a list") | |
| if dtype is not None: | |
| for x in arr: | |
| if not isinstance(x, dtype): | |
| raise ValueError(f"found type {type(x)} in {name}, expect all to be {dtype}") | |
| def linear_blend_skinning( | |
| vertices: ndarray, | |
| matrix_local: ndarray, | |
| matrix: ndarray, | |
| skin: ndarray, | |
| pad: int=1, | |
| value: float=1.0, | |
| ) -> ndarray: | |
| """ | |
| Args: | |
| vertices: (N, 4-pad) | |
| matrix_local: (J, 4, 4) | |
| matrix: (J, 4, 4) | |
| skin: (N, J) | |
| pad: 0 or 1 | |
| value: value to pad | |
| Returns: | |
| (N, 3) vertices using LBS algorithm: Skinning with dual quaternions, Kavan, 2007 | |
| """ | |
| J = matrix_local.shape[0] | |
| N = vertices.shape[0] | |
| assert_ndarray(vertices, name='vertices', shape=(N, 3)) | |
| assert_ndarray(matrix_local, name="matrix_local", shape=(J, 4, 4)) | |
| assert_ndarray(matrix, name="matrix", shape=(J, 4, 4)) | |
| assert_ndarray(skin, name="skin", shape=(N, J)) | |
| assert vertices.shape[-1] + pad == 4 | |
| # (4, N) | |
| padded = np.pad(vertices, ((0, 0), (0, pad)), 'constant', constant_values=(0, value)).T | |
| # (J, 4, 4) | |
| trans = matrix @ np.linalg.inv(matrix_local) | |
| weighted_per_bone_matrix = [] | |
| # (J, N) | |
| mask = (skin > 0).T | |
| for i in range(J): | |
| offset = np.zeros((4, N), dtype=np.float32) | |
| offset[:, mask[i]] = (trans[i] @ padded[:, mask[i]]) * skin.T[i, mask[i]] | |
| weighted_per_bone_matrix.append(offset) | |
| weighted_per_bone_matrix = np.stack(weighted_per_bone_matrix) | |
| g = np.sum(weighted_per_bone_matrix, axis=0) | |
| final = g[:3, :] / (np.sum(skin, axis=1) + 1e-8) | |
| return final.T | |
| def axis_angle_to_matrix(axis_angle: ndarray) -> ndarray: | |
| """ | |
| Turn axis angle representation to matrix representation. | |
| """ | |
| res = np.pad(scipy.spatial.transform.Rotation.from_rotvec(axis_angle).as_matrix(), ((0, 0), (0, 1), (0, 1)), 'constant', constant_values=((0, 0), (0, 0), (0, 0))) | |
| assert res.ndim == 3 | |
| res[:, -1, -1] = 1 | |
| return res | |
| def sample_surface( | |
| num_samples: int, | |
| vertices: ndarray, | |
| faces: ndarray, | |
| mask: Optional[ndarray]=None, | |
| ) -> Tuple[ndarray, ndarray, ndarray]: | |
| ''' | |
| Randomly pick samples proportional to face area. | |
| See sample_surface: https://github.com/mikedh/trimesh/blob/main/trimesh/sample.py | |
| Args: | |
| mask: (num_faces,), only sample points on the faces where value is True. | |
| Return: | |
| vertex_samples: sampled vertices | |
| original_face_index: on which face is sampled | |
| random_lengths: sampled vectors on face | |
| ''' | |
| original_face_indices = np.arange(len(faces)) | |
| # sample according to mask | |
| if mask is not None: | |
| assert_ndarray(arr=mask, name="mask", shape=(faces.shape[0],)) | |
| original_face_indices = original_face_indices[mask] | |
| faces = faces[mask] | |
| # get face area | |
| offset_0 = vertices[faces[:, 1]] - vertices[faces[:, 0]] | |
| offset_1 = vertices[faces[:, 2]] - vertices[faces[:, 0]] | |
| # TODO: change to correct uniform sampling... | |
| face_weight = np.linalg.norm(np.cross(offset_0, offset_1, axis=-1), axis=-1) | |
| weight_cum = np.cumsum(face_weight, axis=0) | |
| face_pick = np.random.rand(num_samples) * weight_cum[-1] | |
| face_index = np.searchsorted(weight_cum, face_pick) | |
| # face_weight = np.cross(offset_0, offset_1, axis=-1) | |
| # face_weight = (face_weight * face_weight).sum(axis=1) | |
| # weight_cum = np.cumsum(face_weight, axis=0) | |
| # face_pick = np.random.rand(num_samples) * weight_cum[-1] | |
| # face_index = np.searchsorted(weight_cum, face_pick) | |
| # map face_index back to original indices | |
| original_face_index = original_face_indices[face_index] | |
| # pull triangles into the form of an origin + 2 vectors | |
| tri_origins = vertices[faces[:, 0]] | |
| tri_vectors = vertices[faces[:, 1:]] | |
| tri_vectors -= np.tile(tri_origins, (1, 2)).reshape((-1, 2, 3)) | |
| # pull the vectors for the faces we are going to sample from | |
| tri_origins = tri_origins[face_index] | |
| tri_vectors = tri_vectors[face_index] | |
| # randomly generate two 0-1 scalar components to multiply edge vectors b | |
| random_lengths = np.random.rand(len(tri_vectors), 2, 1) | |
| random_test = random_lengths.sum(axis=1).reshape(-1) > 1.0 | |
| random_lengths[random_test] -= 1.0 | |
| random_lengths = np.abs(random_lengths) | |
| sample_vector = (tri_vectors * random_lengths).sum(axis=1) | |
| vertex_samples = sample_vector + tri_origins | |
| return vertex_samples, original_face_index, random_lengths | |
| def sample_barycentric( | |
| vertex_group: ndarray, | |
| faces: ndarray, | |
| face_index: ndarray, | |
| random_lengths: ndarray, | |
| ) -> ndarray: | |
| v_origins = vertex_group[faces[face_index, 0]] | |
| v_vectors = vertex_group[faces[face_index, 1:]] | |
| v_vectors -= v_origins[:, np.newaxis, :] | |
| sample_vector = (v_vectors * random_lengths).sum(axis=1) | |
| v_samples = sample_vector + v_origins | |
| return v_samples | |
| def sample_vertex_groups( | |
| vertices: ndarray, | |
| faces: ndarray, | |
| num_samples: int, | |
| num_vertex_samples: Optional[int]=None, | |
| vertex_normals: Optional[ndarray]=None, | |
| face_normals: Optional[ndarray]=None, | |
| vertex_groups: Optional[ndarray]=None, | |
| face_mask: Optional[ndarray]=None, | |
| shuffle: bool=True, | |
| same: bool=False, | |
| ) -> Tuple[ndarray, ndarray|None, ndarray|None]: | |
| """ | |
| Choose num_samples samples on the mesh and get their positions and normals. | |
| If vertex_group is provided, get its weights using barycentric sampling. | |
| Return: | |
| sampled_vertices, sampled_normals, sampled_vertex_groups | |
| Args: | |
| vertices: (N, 3) | |
| faces: (F, 3) | |
| num_samples: how many samples | |
| num_vertex_samples: | |
| At most num_vertex_samples unique vertices to be included, | |
| these points will be concatenated in the last (if shuffle is False). | |
| vertex_normals: (N, 3), sampled_normals will be None if not provided | |
| face_normals: (N, 3), sampled_normals will be None if not provided | |
| vertex_groups: (N, m), sampled_vertex_groups will be None if not provided | |
| face_mask: | |
| (F,) or (F, m), if shape is (F,), use the same mask across all | |
| vertex groups. Only sample on faces where value is True. | |
| shuffle: shuffle samples in the end | |
| same: | |
| Sample on the same locations, only useful when using mutiple | |
| vertex groups and mask is None or shape of (F,). | |
| """ | |
| if num_vertex_samples is None: | |
| num_vertex_samples = 0 | |
| if num_vertex_samples > num_samples: | |
| raise ValueError(f"num_vertex_samples cannot be larger than num_samples, found: {num_vertex_samples} > {num_samples}") | |
| def get_mask_perm(mask: Optional[ndarray]): | |
| if mask is None: | |
| vertex_mask = np.arange(vertices.shape[0]) | |
| else: | |
| vertex_mask = np.unique(mask) | |
| perm = np.random.permutation(vertex_mask.shape[0]) | |
| return vertex_mask[perm[:num_vertex_samples]] | |
| if vertex_groups is not None: | |
| if vertex_groups.ndim == 1: | |
| assert_ndarray(arr=vertex_groups, name="vertex_groups", shape=(vertices.shape[0],)) | |
| vertex_groups = vertex_groups[:, None] | |
| else: | |
| assert_ndarray(arr=vertex_groups, name="vertex_groups", shape=(vertices.shape[0], -1)) | |
| vertex_groups = vertex_groups | |
| if vertex_groups is not None: | |
| if face_mask is not None: | |
| if face_mask.ndim == 1: | |
| assert_ndarray(arr=face_mask, name="mask", shape=(faces.shape[0],)) | |
| else: | |
| assert_ndarray(arr=face_mask, name="mask", shape=(faces.shape[0], vertex_groups.shape[1])) | |
| list_sampled_vertices = [] | |
| list_sampled_normals = [] | |
| list_sampled_vertex_groups = [] | |
| perm = None | |
| _mask = None | |
| same = same and (face_mask is None or (face_mask is not None and face_mask.ndim != 2)) | |
| for i in range(vertex_groups.shape[1]): | |
| if face_mask is not None: | |
| if face_mask.ndim == 1: | |
| perm = get_mask_perm(faces[face_mask]) | |
| _mask = face_mask | |
| else: | |
| perm = get_mask_perm(faces[face_mask[:, i]]) | |
| _mask = face_mask[:, i] | |
| else: | |
| perm = get_mask_perm(None) | |
| _mask = None | |
| _num_samples = num_samples - len(perm) | |
| face_vertices, face_index, random_lengths = sample_surface( | |
| num_samples=_num_samples, | |
| vertices=vertices, | |
| faces=faces, | |
| mask=_mask, | |
| ) | |
| list_sampled_vertices.append(np.concatenate([vertices[perm], face_vertices], axis=0)) | |
| if vertex_normals is not None and face_normals is not None: | |
| list_sampled_normals.append(np.concatenate([vertex_normals[perm], face_normals[face_index]], axis=0)) | |
| if same: | |
| g = sample_barycentric( | |
| vertex_group=vertex_groups, | |
| faces=faces, | |
| face_index=face_index, | |
| random_lengths=random_lengths, | |
| ) | |
| list_sampled_vertex_groups.append(np.concatenate([vertex_groups[perm], g], axis=0)) | |
| break | |
| g = sample_barycentric( | |
| vertex_group=vertex_groups[:, i:i+1], | |
| faces=faces, | |
| face_index=face_index, | |
| random_lengths=random_lengths, | |
| )[:, 0] | |
| list_sampled_vertex_groups.append(np.concatenate([vertex_groups[:, i][perm], g], axis=0)) | |
| sampled_vertices = np.stack(list_sampled_vertices, axis=1) | |
| if len(list_sampled_normals) > 0: | |
| sampled_normals = np.stack(list_sampled_normals, axis=1) | |
| else: | |
| sampled_normals = None | |
| if same: | |
| sampled_vertex_groups = list_sampled_vertex_groups[0] | |
| else: | |
| sampled_vertex_groups = np.stack(list_sampled_vertex_groups, axis=1) | |
| else: # otherwise only sample vertices and normals | |
| if face_mask is not None: | |
| assert_ndarray(arr=face_mask, name="mask", shape=(faces.shape[0],)) | |
| perm = get_mask_perm(faces[face_mask]) | |
| else: | |
| perm = get_mask_perm(None) | |
| num_samples -= len(perm) | |
| n_vertex = vertices[perm] | |
| face_vertices, face_index, random_lengths = sample_surface( | |
| num_samples=num_samples, | |
| vertices=vertices, | |
| faces=faces, | |
| mask=face_mask, | |
| ) | |
| sampled_vertices = np.concatenate([n_vertex, face_vertices], axis=0) | |
| if vertex_normals is not None and face_normals is not None: | |
| sampled_normals = np.concatenate([vertex_normals[perm], face_normals[face_index]], axis=0) | |
| else: | |
| sampled_normals = None | |
| sampled_vertex_groups = None | |
| return sampled_vertices, sampled_normals, sampled_vertex_groups |