GNM Visualization: Clean up type annotations, docstrings, and np.newaxis indexing.

gnm/shape/visualization/gnm_mitsuba.py

@@ -349,7 +349,7 @@ def _prepare_tile_textures(
             interpolation=cv2.INTER_AREA,
         )
       if texture_float.ndim == 2:
-        texture_float = texture_float[..., None]
+        texture_float = texture_float[..., np.newaxis]
       elif texture_float.ndim == 3 and texture_float.shape[2] == 4:
         texture_float = texture_float[..., :3]
       texture_height, texture_width = texture_float.shape[:2]
@@ -424,9 +424,9 @@ def _create_mitsuba_sensor(
 
 
 def _index_optional(
-    array: npt.NDArray[Any] | None,
+    array: npt.NDArray[np.float32] | None,
     index: int | slice | npt.NDArray[np.int32] | tuple[int | slice, ...],
-) -> npt.NDArray[Any] | None:
+) -> npt.NDArray[np.float32] | None:
   """Returns `array[index]` when `array` is not `None`, otherwise `None`."""
   return array[index] if array is not None else None
 
@@ -685,7 +685,7 @@ def render(
     include_shading: bool = True,
     verbose: bool = False,
 ) -> npt.NDArray[np.float32]:
-  """Render GNM meshes using Mitsuba.
+  """Renders GNM meshes using Mitsuba.
 
   N frames, M meshes. See `render_common.BackendRenderFn`.
 
@@ -781,9 +781,11 @@ def render(
     scene_parameters = mi.traverse(scene)
 
     composited_frames = []
-    tqdm_kwargs = dict(
-        disable=not verbose, desc='Rendering (Mitsuba)', leave=False
-    )
+    tqdm_kwargs = {
+        'disable': not verbose,
+        'desc': 'Rendering (Mitsuba)',
+        'leave': False,
+    }
 
     for frame in tqdm.trange(num_frames, **tqdm_kwargs):
       if frame > 0:

gnm/shape/visualization/gnm_mitsuba_test.py

@@ -14,8 +14,6 @@
 
 """Tests for GNM Mitsuba visualization."""
 
-from typing import Any
-
 from absl.testing import absltest
 from absl.testing import parameterized
 import drjit as dr
@@ -66,26 +64,27 @@ def setUp(self):
         dtype=np.float32,
     )
 
-  def broadcast(
-      self, array: npt.NDArray[Any], num_frames: int = 1
-  ) -> npt.NDArray[Any]:
-    return np.broadcast_to(array, (num_frames, *array.shape))
-
   @parameterized.product(
       version=_MAINTAINED_MAJOR_GNM_VERSIONS,
       num_frames=(1, 3),
   )
   def test_render_basic(self, version: str, num_frames: int):
     """Tests basic rendering with Mitsuba."""
     gnm_np = self.gnms[version]
-    vertices = gnm_np.template_vertex_positions[None, :, :]
-    vertices = self.broadcast(vertices, num_frames)
+    vertices = np.broadcast_to(
+        gnm_np.template_vertex_positions,
+        (num_frames, 1, *gnm_np.template_vertex_positions.shape),
+    )
 
     color = gnm_mitsuba.render(
         vertices=vertices,
         triangles=gnm_np.triangles,
-        world_to_camera=self.broadcast(self.world_to_camera, num_frames),
-        camera_to_image=self.broadcast(self.camera_to_image, num_frames),
+        world_to_camera=np.broadcast_to(
+            self.world_to_camera, (num_frames, 4, 4)
+        ),
+        camera_to_image=np.broadcast_to(
+            self.camera_to_image, (num_frames, 4, 4)
+        ),
         vertex_normals=gnm_np.compute_vertex_normals(vertices),
         triangle_uvs=gnm_np.triangle_uvs_udim,
         vertex_colors=np.ones_like(vertices),
@@ -178,16 +177,16 @@ def test_create_mitsuba_mesh_no_colors(self):
   def test_render_partial_texture_dict(self):
     """Tests that missing tiles in texture dict fall back to default texture."""
     gnm_np = self.gnms[_MAINTAINED_MAJOR_GNM_VERSIONS[0]]
-    vertices = gnm_np.template_vertex_positions[None, None, ...]
+    vertices = gnm_np.template_vertex_positions[np.newaxis, np.newaxis]
     # Only provide a texture for the skin tile, omitting the other tiles.
-    partial_texture: dict[int, np.ndarray] = {
+    partial_texture: dict[int, npt.NDArray[np.float32]] = {
         gnm_udim.UdimTile.SKIN: np.ones((1, 16, 16, 3), np.float32),
     }
     color = gnm_mitsuba.render(
         vertices=vertices,
         triangles=gnm_np.triangles,
-        world_to_camera=self.broadcast(self.world_to_camera, 1),
-        camera_to_image=self.broadcast(self.camera_to_image, 1),
+        world_to_camera=self.world_to_camera[np.newaxis],
+        camera_to_image=self.camera_to_image[np.newaxis],
         vertex_normals=gnm_np.compute_vertex_normals(vertices),
         triangle_uvs=gnm_np.triangle_uvs_udim,
         vertex_colors=np.ones_like(vertices),
@@ -203,7 +202,7 @@ def test_render_multi_frame_updates_in_place(self):
     num_frames = 3
     vertices = np.stack(
         [
-            base_vertices[None, ...]
+            base_vertices[np.newaxis]
             + np.array(
                 [0.02 * frame, -0.01 * frame, 0.0],
                 dtype=np.float32,
@@ -245,7 +244,7 @@ def test_render_multi_frame_updates_in_place(self):
 
     skin_texture_first = np.full((8, 8, 3), 0.8, dtype=np.float32)
     skin_texture_last = np.full((8, 8, 3), 0.5, dtype=np.float32)
-    texture: dict[int, np.ndarray] = {
+    texture: dict[int, npt.NDArray[np.float32]] = {
         # Frame 0 and Frame 1 share the same texture; Frame 2 changes texture.
         gnm_udim.UdimTile.SKIN: np.stack(
             [skin_texture_first, skin_texture_first, skin_texture_last], axis=0
@@ -288,16 +287,16 @@ def test_render_multi_frame_updates_in_place(self):
   def test_non_square_pixels_raises_error(self):
     """Tests that non-square pixels raise an error."""
     gnm_np = self.gnms[_MAINTAINED_MAJOR_GNM_VERSIONS[0]]
-    vertices = gnm_np.template_vertex_positions[None, None, ...]
+    vertices = gnm_np.template_vertex_positions[np.newaxis, np.newaxis]
     non_square_c2i = self.camera_to_image.copy()
     non_square_c2i[0, 0] = 500.0
     non_square_c2i[1, 1] = 600.0
     with self.assertRaisesRegex(NotImplementedError, 'square pixels'):
       gnm_mitsuba.render(
           vertices=vertices,
           triangles=gnm_np.triangles,
-          world_to_camera=self.broadcast(self.world_to_camera, 1),
-          camera_to_image=self.broadcast(non_square_c2i, 1),
+          world_to_camera=self.world_to_camera[np.newaxis],
+          camera_to_image=non_square_c2i[np.newaxis],
           vertex_normals=gnm_np.compute_vertex_normals(vertices),
           triangle_uvs=gnm_np.triangle_uvs_udim,
           vertex_colors=np.ones_like(vertices),
@@ -307,15 +306,15 @@ def test_non_square_pixels_raises_error(self):
   def test_skew_raises_error(self):
     """Tests that skew raises an error."""
     gnm_np = self.gnms[_MAINTAINED_MAJOR_GNM_VERSIONS[0]]
-    vertices = gnm_np.template_vertex_positions[None, None, ...]
+    vertices = gnm_np.template_vertex_positions[np.newaxis, np.newaxis]
     skew_c2i = self.camera_to_image.copy()
     skew_c2i[0, 1] = 10.0
     with self.assertRaisesRegex(NotImplementedError, 'zero skew'):
       gnm_mitsuba.render(
           vertices=vertices,
           triangles=gnm_np.triangles,
-          world_to_camera=self.broadcast(self.world_to_camera, 1),
-          camera_to_image=self.broadcast(skew_c2i, 1),
+          world_to_camera=self.world_to_camera[np.newaxis],
+          camera_to_image=skew_c2i[np.newaxis],
           vertex_normals=gnm_np.compute_vertex_normals(vertices),
           triangle_uvs=gnm_np.triangle_uvs_udim,
           vertex_colors=np.ones_like(vertices),
@@ -325,7 +324,7 @@ def test_skew_raises_error(self):
   def test_render_suspends_gradients(self):
     """Tests that mesh creation and rendering suspend gradient tracking."""
     gnm_np = self.gnms[_MAINTAINED_MAJOR_GNM_VERSIONS[0]]
-    vertices = gnm_np.template_vertex_positions[None, None, ...]
+    vertices = gnm_np.template_vertex_positions[np.newaxis, np.newaxis]
     tracked_value = mi.Float(1.0)
     dr.enable_grad(tracked_value)
     with dr.resume_grad():
@@ -340,8 +339,8 @@ def test_render_suspends_gradients(self):
       color = gnm_mitsuba.render(
           vertices=vertices,
           triangles=gnm_np.triangles,
-          world_to_camera=self.broadcast(self.world_to_camera, 1),
-          camera_to_image=self.broadcast(self.camera_to_image, 1),
+          world_to_camera=self.world_to_camera[np.newaxis],
+          camera_to_image=self.camera_to_image[np.newaxis],
           vertex_normals=gnm_np.compute_vertex_normals(vertices),
           triangle_uvs=gnm_np.triangle_uvs_udim,
           vertex_colors=np.ones_like(vertices),

gnm/shape/visualization/gnm_pyrender.py

@@ -18,14 +18,16 @@
 arguments; call `render_gnm.render_gnm` rather than `render` directly.
 """
 
-import collections.abc
+from collections.abc import Callable, Mapping
 import functools
 import os
 import tempfile
 from typing import Any
+
 import cv2
 from gnm.shape.visualization import udim
 import numpy as np
+import numpy.typing as npt
 import tqdm
 
 if 'PYOPENGL_PLATFORM' not in os.environ:
@@ -46,21 +48,26 @@
 class ProjectionMatrixCamera(pyrender.Camera):
   """A camera defined by a full projection matrix."""
 
-  def __init__(self, projection_matrix: np.ndarray):
+  def __init__(self, projection_matrix: npt.NDArray[np.floating]):
+    """Initializes the camera with the given 4x4 projection matrix."""
     super().__init__()
     self._projection_matrix = projection_matrix
 
   def get_projection_matrix(
       self, width: int | None = None, height: int | None = None
-  ) -> np.ndarray:
+  ) -> npt.NDArray[np.floating]:
+    """Returns the 4x4 projection matrix."""
     del width, height
     return self._projection_matrix
 
-  def set_projection_matrix(self, projection_matrix: np.ndarray):
+  def set_projection_matrix(
+      self, projection_matrix: npt.NDArray[np.floating]
+  ) -> None:
+    """Updates the 4x4 projection matrix."""
     self._projection_matrix = projection_matrix
 
 
-def _changed_flag(array: np.ndarray) -> np.ndarray:
+def _changed_flag(array: npt.NDArray[np.generic]) -> npt.NDArray[np.bool_]:
   """Returns a flag indicating whether the array has changed between frames.
 
   Args:
@@ -70,7 +77,6 @@ def _changed_flag(array: np.ndarray) -> np.ndarray:
     A flag indicating whether the array has changed between frames, (N,). Always
     1 for the first frame.
   """
-
   if array.strides[0] == 0:
     # 'N' dimension is just a broadcast.
     arr = np.zeros(array.shape[0], dtype=np.bool_)
@@ -82,25 +88,27 @@ def _changed_flag(array: np.ndarray) -> np.ndarray:
 
 
 def render(
-    vertices: np.ndarray,
-    triangles: np.ndarray,
-    world_to_camera: np.ndarray,
-    camera_to_image: np.ndarray,
-    vertex_normals: np.ndarray,
-    triangle_uvs: np.ndarray,
-    vertex_colors: np.ndarray,
+    vertices: npt.NDArray[np.floating],
+    triangles: npt.NDArray[np.integer],
+    world_to_camera: npt.NDArray[np.floating],
+    camera_to_image: npt.NDArray[np.floating],
+    vertex_normals: npt.NDArray[np.floating],
+    triangle_uvs: npt.NDArray[np.floating],
+    vertex_colors: npt.NDArray[np.floating],
     image_size: tuple[int, int] = (240, 320),
-    texture: collections.abc.Mapping[int, np.ndarray] | None = None,
+    texture: (
+        Mapping[int, npt.NDArray[np.floating] | npt.NDArray[np.uint8]] | None
+    ) = None,
     multisample_antialiasing: int = 1,
-    background_color: np.ndarray | None = None,
+    background_color: npt.NDArray[np.floating] | None = None,
     alpha: float = 1.0,
     include_shading: bool = True,
     verbose: bool = False,
-    renderer_factory: collections.abc.Callable[
+    renderer_factory: Callable[
         [int, int], pyrender.OffscreenRenderer
     ] = DEFAULT_RENDERER_FACTORY,
-) -> np.ndarray:
-  """Render GNM meshes.
+) -> npt.NDArray[np.float32]:
+  """Renders GNM meshes.
 
   N frames, M meshes. See `render_common.BackendRenderFn`.
 
@@ -127,9 +135,8 @@ def render(
       pyrender.OffscreenRenderer.
 
   Returns:
-    The rendered color image, float32 [0-1] (H, W, 3).
+    The rendered color image, float32 [0-1] (N, H, W, 3).
   """
-
   num_frames, num_meshes = vertices.shape[:2]
 
   # Use multi-sample anti-aliasing (MSAA) to reduce jagged edges.
@@ -182,7 +189,7 @@ def _apply_texture(
       mesh_node: pyrender.Node,
       texture: pyrender.Texture | None,
       tile_index: int,
-  ):
+  ) -> None:
     material = mesh_node.mesh.primitives[tile_index].material
     material.baseColorTexture = texture
 
@@ -220,7 +227,7 @@ def _apply_texture(
 
   texture_objects = {}
   colors = []
-  tqdm_kwargs = dict(disable=not verbose, desc='Rendering', leave=False)
+  tqdm_kwargs = {'disable': not verbose, 'desc': 'Rendering', 'leave': False}
   for f in tqdm.trange(0, num_frames, **tqdm_kwargs):
     # Update anything that's changed.
     if vertices_changed[f] or vertex_colors_changed[f]:
@@ -267,8 +274,8 @@ def _apply_texture(
     color = np.power(color, 2.2)
 
     # Alpha blending with background color.
-    color = color * alpha_im[..., None] + background * (
-        1.0 - alpha_im[..., None]
+    color = color * alpha_im[..., np.newaxis] + background * (
+        1.0 - alpha_im[..., np.newaxis]
     )
 
     # Downsample for anti-aliasing.
@@ -278,7 +285,7 @@ def _apply_texture(
     colors.append(color)
 
   renderer.delete()
-  return np.array(colors)
+  return np.array(colors, dtype=np.float32)
 
 
 @functools.cache
@@ -311,6 +318,7 @@ class CustomShaderCache(pyrender.shader_program.ShaderProgramCache):
   def __init__(
       self, original_cache: pyrender.shader_program.ShaderProgramCache
   ):
+    """Initializes the cache wrapper around `original_cache`."""
     super().__init__()
     self.original_cache = original_cache
     self.shader_program = None
@@ -324,7 +332,7 @@ def get_program(
       *args,
       **kwargs,
   ) -> pyrender.shader_program.ShaderProgram:
-
+    """Returns a shader program with a custom fragment shader."""
     # Remove sRGB correction from textures - our textures are in linear space.
     defines = defines or {}
     defines['SRGB_CORRECTED'] = 0
@@ -350,17 +358,18 @@ def get_program(
     return program
 
 
-def _override_shader_cache(renderer: pyrender.OffscreenRenderer):
-  """Override the shader cache with a custom cache."""
-
+def _override_shader_cache(renderer: pyrender.OffscreenRenderer) -> None:
+  """Overrides the shader cache with a custom cache."""
   # pylint: disable=protected-access
   if r := renderer._renderer:
     cache = r._program_cache
     r._program_cache = CustomShaderCache(cache)
   # pylint: enable=protected-access
 
 
-def _align_vectors(a: np.ndarray, b: np.ndarray) -> np.ndarray:
+def _align_vectors(
+    a: npt.NDArray[np.floating], b: npt.NDArray[np.floating]
+) -> npt.NDArray[np.float32]:
   """Returns a 4x4 transform that rotates vector `a` to vector `b`."""
   a = np.asarray(a, dtype=np.float64)
   b = np.asarray(b, dtype=np.float64)