From fd3eace0ba47e2162018a2f80fb3d161808800c8 Mon Sep 17 00:00:00 2001 From: manumishra12 Date: Thu, 16 Jul 2026 08:31:56 -0700 Subject: [PATCH] Improve docstring and README consistency in gnm.shape - gnm_numpy/gnm_jax/gnm_tensorflow: give compute_vertex_normals the same Args/Returns docstring the PyTorch backend already has. - gnm_tensorflow: document the GNM attributes that the other backends list but the TF class docstring omitted (quad_uvs, triangle_uvs, mesh_component_names, mirror_indices, variant, vertex_groups, vertex_group_names, joint_connections, joint_children_indices, skinning_segmentation, edge_list, vertex_uvs) and add the missing P shape dimension. - README: the per-joint rotations are axis-angle vectors of shape [batch_size, num_joints, 3], not a "4x3 Rotation matrix". --- gnm/shape/README.md | 2 +- gnm/shape/gnm_jax.py | 9 ++++++++- gnm/shape/gnm_numpy.py | 9 ++++++++- gnm/shape/gnm_tensorflow.py | 25 ++++++++++++++++++++++++- 4 files changed, 41 insertions(+), 4 deletions(-) diff --git a/gnm/shape/README.md b/gnm/shape/README.md index 00e0562d..eef774f5 100644 --- a/gnm/shape/README.md +++ b/gnm/shape/README.md @@ -230,7 +230,7 @@ relevant for the GNM v3.x. ### Joint Parameters -* **Shape:** rotations: `[batch_size, 4x3 Rotation matrix]`, global translation: `[batch_size, 3]` +* **Shape:** rotations: `[batch_size, num_joints, 3]` (axis-angle), global translation: `[batch_size, 3]` * **Description:** Controls the global head position and joint angles for head pose and eyeball orientation. ## Model Data diff --git a/gnm/shape/gnm_jax.py b/gnm/shape/gnm_jax.py index 88c589c3..52e56bd9 100644 --- a/gnm/shape/gnm_jax.py +++ b/gnm/shape/gnm_jax.py @@ -125,7 +125,14 @@ def compute_vertex_normals( self, vertices: jt.Float[jnp.ndarray, "... V 3"], ) -> jt.Float[jnp.ndarray, "... V 3"]: - """Compute vertex normals for GNM mesh.""" + """Compute vertex normals for GNM mesh. + + Args: + vertices: Vertex positions with shape `(..., V, 3)`. + + Returns: + Vertex normals with shape `(..., V, 3)`. + """ batch_dims = vertices.shape[:-2] num_vertices = vertices.shape[-2] vertices_flat = vertices.reshape(-1, num_vertices, 3) diff --git a/gnm/shape/gnm_numpy.py b/gnm/shape/gnm_numpy.py index 37340d12..0033e716 100644 --- a/gnm/shape/gnm_numpy.py +++ b/gnm/shape/gnm_numpy.py @@ -124,7 +124,14 @@ def compute_vertex_normals( self, vertices: npt.NDArray[np.floating], ) -> npt.NDArray[np.floating]: - """Compute vertex normals for GNM mesh.""" + """Compute vertex normals for GNM mesh. + + Args: + vertices: Vertex positions with shape `(..., V, 3)`. + + Returns: + Vertex normals with shape `(..., V, 3)`. + """ batch_dims = vertices.shape[:-2] num_vertices = vertices.shape[-2] vertices_flat = vertices.reshape(-1, num_vertices, 3) diff --git a/gnm/shape/gnm_tensorflow.py b/gnm/shape/gnm_tensorflow.py index 059f2dc4..da4e546d 100644 --- a/gnm/shape/gnm_tensorflow.py +++ b/gnm/shape/gnm_tensorflow.py @@ -74,6 +74,7 @@ class GNM(gnm_xnp.GNM): * Q: The number of quads in the mesh topology. * T: The number of triangles, in a triangulated version of the mesh topology. * G: Number of vertex groups. + * P: Number of separate parts of the mesh. Attributes: joint_parent_indices: Parent's index for each joint in the skeleton, (J). @@ -95,6 +96,21 @@ class GNM(gnm_xnp.GNM): 3). If they do not exist in the GNM npz, they are set to the identity matrix. Note that these are not used to compute the GNM joint and vertex positions. + quad_uvs: Texture coordinates per quad, (Q, 4, 2). + triangle_uvs: Texture coordinates per triangle, (T, 3, 2). + mesh_component_names: The vertex group name corresponding to each separate + mesh part. + mirror_indices: The index of each vertex on the other side of the mesh. + variant: The variant of the loaded GNM model. + vertex_groups: The weights in each vertex group, (G, V). + vertex_group_names: The name of each vertex group, (G,). + joint_connections: The joint connections of the GNM rig. + joint_children_indices: A dictionary that contains the joint indices of the + children of each joint. + skinning_segmentation: A decomposition of the vertices into separate parts + based on skinning weights. + edge_list: The quad topology represented as a list of directed edges (E, 2). + vertex_uvs: Per-vertex UV texture coordinates shaped (V, 2). num_vertices: The number of vertices in the mesh V. num_joints: The number of joints in the skeleton J. identity_dim: The dimensionality of the linear identity basis I. @@ -113,7 +129,14 @@ def compute_vertex_normals( self, vertices: tf.Tensor, ) -> tf.Tensor: - """Compute vertex normals for GNM mesh.""" + """Compute vertex normals for GNM mesh. + + Args: + vertices: Vertex positions with shape `(..., V, 3)`. + + Returns: + Vertex normals with shape `(..., V, 3)`. + """ batch_dims = tf.shape(vertices)[:-2] num_vertices = tf.shape(vertices)[-2] vertices_flat = tf.reshape(vertices, (-1, num_vertices, 3))