From f0dd7d1774f72b16adf50f9c5dbfc8eadbe6c91a Mon Sep 17 00:00:00 2001 From: Sebastian Winberg Date: Wed, 28 Oct 2026 10:40:58 -0700 Subject: [PATCH] GNM Visualization: Clean up type annotations, docstrings, and `np.newaxis` indexing. - Refine array type annotations (`npt.NDArray[...]`, `FloatArray`) and return annotations across `gnm_mitsuba`, `gnm_pyrender`, `render_common`, and their unit tests. - Standardize docstrings, replace `dict(...)` calls with dict literals, and use `np.newaxis` instead of `None` for array axis expansion. PiperOrigin-RevId: 1009020035 --- gnm/shape/visualization/gnm_mitsuba.py | 16 +- gnm/shape/visualization/gnm_mitsuba_test.py | 51 +++--- gnm/shape/visualization/gnm_pyrender.py | 69 ++++---- gnm/shape/visualization/gnm_pyrender_test.py | 32 ++-- gnm/shape/visualization/render_common.py | 158 ++++++++++-------- gnm/shape/visualization/render_common_test.py | 17 +- 6 files changed, 189 insertions(+), 154 deletions(-) diff --git a/gnm/shape/visualization/gnm_mitsuba.py b/gnm/shape/visualization/gnm_mitsuba.py index 9168df09..8ebecc89 100644 --- a/gnm/shape/visualization/gnm_mitsuba.py +++ b/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: diff --git a/gnm/shape/visualization/gnm_mitsuba_test.py b/gnm/shape/visualization/gnm_mitsuba_test.py index 25d4d4b8..993090fb 100644 --- a/gnm/shape/visualization/gnm_mitsuba_test.py +++ b/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,11 +64,6 @@ 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), @@ -78,14 +71,20 @@ def broadcast( 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,7 +287,7 @@ 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 @@ -296,8 +295,8 @@ def test_non_square_pixels_raises_error(self): 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), diff --git a/gnm/shape/visualization/gnm_pyrender.py b/gnm/shape/visualization/gnm_pyrender.py index cf56b390..5d646620 100644 --- a/gnm/shape/visualization/gnm_pyrender.py +++ b/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,9 +358,8 @@ 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 @@ -360,7 +367,9 @@ def _override_shader_cache(renderer: pyrender.OffscreenRenderer): # 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) diff --git a/gnm/shape/visualization/gnm_pyrender_test.py b/gnm/shape/visualization/gnm_pyrender_test.py index a9cae6ce..2849ab40 100644 --- a/gnm/shape/visualization/gnm_pyrender_test.py +++ b/gnm/shape/visualization/gnm_pyrender_test.py @@ -14,9 +14,8 @@ """Tests for GNM pyrender visualization.""" -# pylint: disable=protected-access - import tempfile + from absl.testing import absltest from absl.testing import parameterized from etils import epath @@ -61,23 +60,26 @@ def setUp(self): self.camera_to_image[1, 1] = 10 self.camera_to_image[2:4, 2] = -1.0 - def broadcast(self, array: np.ndarray, num_frames: int = 1) -> np.ndarray: - return np.broadcast_to(array, (num_frames, *array.shape)) - @parameterized.product( version=_MAINTAINED_MAJOR_GNM_VERSIONS, num_frames=(1, 5), ) def test_render_basic(self, version: str, num_frames: int): 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_pyrender.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), @@ -121,14 +123,18 @@ def test_render_clips_out_of_range_float_texture(self): class ProjectionMatrixCameraTest(parameterized.TestCase): @parameterized.parameters((np.eye(4),), (np.zeros((4, 4)),)) - def test_projection_matrix_camera(self, projection_matrix: np.ndarray): + def test_projection_matrix_camera( + self, projection_matrix: npt.NDArray[np.floating] + ): camera = gnm_pyrender.ProjectionMatrixCamera(projection_matrix) np.testing.assert_array_equal( camera.get_projection_matrix(), projection_matrix ) @parameterized.parameters((np.eye(4),), (np.zeros((4, 4)),)) - def test_set_projection_matrix(self, projection_matrix: np.ndarray): + def test_set_projection_matrix( + self, projection_matrix: npt.NDArray[np.floating] + ): camera = gnm_pyrender.ProjectionMatrixCamera(np.ones((4, 4))) camera.set_projection_matrix(projection_matrix) np.testing.assert_array_equal( @@ -152,7 +158,7 @@ def test_changed_flag(self): array[index, 0, 0] = 1.0 np.testing.assert_array_equal( - gnm_pyrender._changed_flag(array), expected_flags + gnm_pyrender._changed_flag(array), expected_flags # pylint: disable=protected-access ) def test_changed_flag_broadcast_unchanged(self): @@ -165,7 +171,7 @@ def test_changed_flag_broadcast_unchanged(self): expected_flags[0] = True np.testing.assert_array_equal( - gnm_pyrender._changed_flag(array), expected_flags + gnm_pyrender._changed_flag(array), expected_flags # pylint: disable=protected-access ) diff --git a/gnm/shape/visualization/render_common.py b/gnm/shape/visualization/render_common.py index 23bc1d1f..69c27417 100644 --- a/gnm/shape/visualization/render_common.py +++ b/gnm/shape/visualization/render_common.py @@ -21,7 +21,7 @@ from collections.abc import Mapping, Sequence import functools -from typing import Protocol +from typing import Any, Protocol, overload import warnings from etils import epath @@ -93,20 +93,22 @@ class BackendRenderFn(Protocol): def __call__( self, vertices: FloatArray, - triangles: np.ndarray, + triangles: npt.NDArray[np.integer], world_to_camera: FloatArray, camera_to_image: FloatArray, vertex_normals: FloatArray, triangle_uvs: FloatArray, vertex_colors: FloatArray, image_size: tuple[int, int] = DEFAULT_IMAGE_SIZE, - texture: Mapping[int, np.ndarray] | None = None, + texture: ( + Mapping[int, FloatArray | npt.NDArray[np.uint8]] | None + ) = None, multisample_antialiasing: int = 1, background_color: FloatArray | None = None, alpha: float = 1.0, include_shading: bool = True, verbose: bool = False, - ) -> np.ndarray: + ) -> FloatArray: ... @@ -179,7 +181,7 @@ def render_gnm_mesh( vertices = gnm_np.template_vertex_positions if not multiple_gnms: - vertices = vertices[..., None, :, :] # Inject 'M' dimension. + vertices = vertices[..., np.newaxis, :, :] # Inject 'M' dimension. elif (vertices_dim := vertices.ndim) < 3: raise ValueError( f'Called with {multiple_gnms=}, but vertices is only {vertices_dim}D.' @@ -201,7 +203,7 @@ def render_gnm_mesh( if not colors_has_m_dim: # 'M' dimension is not present but is required. Inject the dimension, and # expand to match the size of vertices' M dimension. - vertex_colors = vertex_colors[..., None, :, :] + vertex_colors = vertex_colors[..., np.newaxis, :, :] vertex_colors_shape = list(vertex_colors.shape) vertex_colors_shape[-3] = vertices.shape[-3] vertex_colors = np.broadcast_to(vertex_colors, vertex_colors_shape) @@ -244,6 +246,8 @@ def render_gnm_mesh( # Convert background color to float [0-1]. if background_color.dtype == np.uint8: background_color = background_color.astype(np.float32) / 255.0 + else: + background_color = background_color.astype(np.float32, copy=False) texture_dict = load_texture(gnm_np, texture) textures = list(texture_dict.values()) @@ -267,8 +271,10 @@ def render_gnm_mesh( f' {[t.shape for t in textures]}.' ) from e - def batchify(arr, non_batch_dims): - """Broadcast to batch dimensions, and flatten the batch dimensions.""" + def batchify[ScalarT: np.generic]( + arr: npt.NDArray[ScalarT], non_batch_dims: int + ) -> npt.NDArray[ScalarT]: + """Broadcasts to batch dimensions, and flattens the batch dimensions.""" arr = np.broadcast_to(arr, (*batch_dims, *arr.shape[-non_batch_dims:])) return arr.reshape(int(np.prod(batch_dims)), *arr.shape[-non_batch_dims:]) @@ -303,14 +309,14 @@ def batchify(arr, non_batch_dims): def project_points_for_gnm( gnm_np: gnm_numpy.GNM, - points_world: np.ndarray | None = None, - vertices: np.ndarray | None = None, + points_world: FloatArray | None = None, + vertices: FloatArray | None = None, world_to_camera: FloatArray | None = None, camera_to_image: FloatArray | None = None, image_size: tuple[int, int] = DEFAULT_IMAGE_SIZE, multiple_gnms: bool = False, - **kwargs, -) -> np.ndarray: + **kwargs: Any, +) -> FloatArray: """Projects world points under the same conditions as a render_gnm call. Intended for identifying the per-frame positions of 3D points (e.g. GNM @@ -336,7 +342,6 @@ def project_points_for_gnm( Returns: The projected points in image space, (..., P, 2). """ - del kwargs if points_world is None: @@ -347,7 +352,7 @@ def project_points_for_gnm( vertices = gnm_np.template_vertex_positions if not multiple_gnms: - vertices = vertices[..., None, :, :] # Inject 'M' dimension. + vertices = vertices[..., np.newaxis, :, :] # Inject 'M' dimension. elif (vertices_dim := vertices.ndim) < 3: raise ValueError( f'Called with {multiple_gnms=}, but vertices is only {vertices_dim}D.' @@ -396,8 +401,8 @@ def project_points_for_gnm( f' camera_to_image {camera_to_image.shape}.' ) from e - def batchify(arr, non_batch_dims): - """Broadcast to batch dimensions, and flatten the batch dimensions.""" + def batchify(arr: FloatArray, non_batch_dims: int) -> FloatArray: + """Broadcasts to batch dimensions, and flattens the batch dimensions.""" arr = np.broadcast_to(arr, (*batch_dims, *arr.shape[-non_batch_dims:])) return arr.reshape(int(np.prod(batch_dims)), *arr.shape[-non_batch_dims:]) @@ -413,7 +418,8 @@ def batchify(arr, non_batch_dims): ) points_homogeneous = np.concatenate([points_world, homogenous_ones], axis=-1) points_clip_space = ( - view_projection_matrix[:, None, :, :] @ points_homogeneous[..., None] + view_projection_matrix[:, np.newaxis] + @ points_homogeneous[..., np.newaxis] ) points_clip_space = points_clip_space[..., 0] @@ -435,18 +441,18 @@ def batchify(arr, non_batch_dims): def get_look_at_world_to_camera( gnm_np: gnm_numpy.GNM, - vertices_world: np.ndarray, - azimuthal_angle: np.ndarray | float = 0.0, - polar_angle: np.ndarray | float = 0.0, - camera_distance: np.ndarray | float | None = DEFAULT_CAMERA_DISTANCE, - share_camera: np.ndarray | bool = True, - y_up: np.ndarray | bool = False, + vertices_world: FloatArray, + azimuthal_angle: FloatArray | float = 0.0, + polar_angle: FloatArray | float = 0.0, + camera_distance: FloatArray | float | None = DEFAULT_CAMERA_DISTANCE, + share_camera: npt.NDArray[np.bool_] | bool = True, + y_up: npt.NDArray[np.bool_] | bool = False, look_at_vertex_groups: Sequence[str] = ('hockey_mask',), left_vertex_groups: Sequence[str] = ('ears', '&left'), right_vertex_groups: Sequence[str] = ('ears', '&right'), forward_vertex_groups: Sequence[str] = ('nose_region',), -) -> np.ndarray: - """Compute world-to-camera matrices for a 'look-at' transform. +) -> FloatArray: + """Computes world-to-camera matrices for a 'look-at' transform. Returns matrices in OpenCV convention. @@ -468,7 +474,6 @@ def get_look_at_world_to_camera( Returns: The world-to-camera matrices, (..., 4, 4). """ - batch_dims = vertices_world.shape[:-2] if not batch_dims: @@ -489,7 +494,7 @@ def get_look_at_world_to_camera( ) vertices_for_camera = np.where( - share_camera[..., None], first_frame_vertices, vertices_world + share_camera[..., np.newaxis], first_frame_vertices, vertices_world ) gnm_axes = _get_gnm_axes( @@ -524,13 +529,13 @@ def get_look_at_world_to_camera( def get_fill_factor_camera_to_image( gnm_np: gnm_numpy.GNM, - vertices: np.ndarray, - target_fill_factor: np.ndarray | float = DEFAULT_TARGET_FILL_FACTOR, - camera_distance: np.ndarray | float = DEFAULT_CAMERA_DISTANCE, + vertices: FloatArray, + target_fill_factor: FloatArray | float = DEFAULT_TARGET_FILL_FACTOR, + camera_distance: FloatArray | float = DEFAULT_CAMERA_DISTANCE, image_size: tuple[int, int] = DEFAULT_IMAGE_SIZE, vertex_group_name: str = 'hockey_mask', -) -> np.ndarray: - """Compute the camera-to-image matrix for a 'fill factor' transform. +) -> FloatArray: + """Computes the camera-to-image matrix for a 'fill factor' transform. The camera intrinsics are determined to fill the width of the given vertex group with the image. Returns matrices in OpenCV convention. @@ -581,15 +586,15 @@ def get_fill_factor_camera_to_image( def get_spin_world_to_camera( gnm_np: gnm_numpy.GNM, - vertices: np.ndarray, + vertices: FloatArray, has_time_dimension: bool = False, spin_period: int = 60, spin_azimuth_limit: float = 20.0, spin_polar_limit: float = 5.0, num_frames: int | None = None, - **kwargs, -) -> np.ndarray: - """Compute world-to-camera matrices for a 'spin' transform. + **kwargs: Any, +) -> FloatArray: + """Computes world-to-camera matrices for a 'spin' transform. Args: gnm_np: The GNM model. @@ -606,7 +611,6 @@ def get_spin_world_to_camera( The world-to-camera matrices, (spin_period, ..., 4, 4) if has_time_dimension is False, otherwise (..., 4, 4). """ - if num_frames is not None: if not has_time_dimension and ( vertices.ndim < 3 or vertices.shape[0] != num_frames @@ -624,8 +628,8 @@ def get_spin_world_to_camera( max_angle = np.pi * 2 * (num_frames / spin_period) angles = np.linspace(-0, max_angle, num_frames) - azimuthal_angle = np.sin(angles)[..., None] * spin_azimuth_limit - polar_angle = np.cos(angles)[..., None] * spin_polar_limit + azimuthal_angle = np.sin(angles)[..., np.newaxis] * spin_azimuth_limit + polar_angle = np.cos(angles)[..., np.newaxis] * spin_polar_limit # Broadcast azimuthal and polar angles to the batch dimensions. azimuthal_angle = np.broadcast_to(azimuthal_angle, (*batch_dims, 1)) @@ -641,12 +645,12 @@ def get_spin_world_to_camera( def _get_gnm_axes( - vertices: np.ndarray, + vertices: FloatArray, left_vertex_groups: Sequence[str], right_vertex_groups: Sequence[str], forward_vertex_groups: Sequence[str], gnm_np: gnm_numpy.GNM, -) -> tuple[np.ndarray, np.ndarray, np.ndarray]: +) -> tuple[FloatArray, FloatArray, FloatArray]: """Determines the right, up, and forwards direction vectors of GNM. Args: @@ -659,7 +663,6 @@ def _get_gnm_axes( Returns: A tuple of right, up, and forwards direction vectors, each (..., 3). """ - left = _vertex_group_mean(vertices, left_vertex_groups, gnm_np) right = _vertex_group_mean(vertices, right_vertex_groups, gnm_np) forward_point = _vertex_group_mean(vertices, forward_vertex_groups, gnm_np) @@ -673,8 +676,8 @@ def _get_gnm_axes( def right_handed_look_at( - camera_position: np.ndarray, look_at: np.ndarray, up_vector: np.ndarray -) -> np.ndarray: + camera_position: FloatArray, look_at: FloatArray, up_vector: FloatArray +) -> FloatArray: """Builds a right handed look at view matrix. Args: @@ -685,7 +688,6 @@ def right_handed_look_at( Returns: The right handed look at view matrix, (..., 4, 4). """ - z_axis = look_at - camera_position z_axis /= np.linalg.norm(z_axis, axis=-1, keepdims=True) @@ -694,8 +696,8 @@ def right_handed_look_at( vertical_axis = np.cross(horizontal_axis, z_axis, axis=-1) - def _dot_last_axis(arr1, arr2): - return np.einsum('...i,...i->...', arr1, arr2)[..., None] + def _dot_last_axis(arr1: FloatArray, arr2: FloatArray) -> FloatArray: + return np.einsum('...i,...i->...', arr1, arr2)[..., np.newaxis] batch_shape = horizontal_axis.shape[:-1] zeros = np.zeros((*batch_shape, 3), dtype=horizontal_axis.dtype) @@ -714,11 +716,11 @@ def _dot_last_axis(arr1, arr2): def _get_camera_offset( - head_axes: tuple[np.ndarray, np.ndarray, np.ndarray], - camera_distance: np.ndarray, - azimuthal_angle: np.ndarray, - polar_angle: np.ndarray, -) -> np.ndarray: + head_axes: tuple[FloatArray, FloatArray, FloatArray], + camera_distance: FloatArray, + azimuthal_angle: FloatArray, + polar_angle: FloatArray, +) -> FloatArray: """Determines how to place the camera so it points at the face. Args: @@ -741,11 +743,11 @@ def _get_camera_offset( def _right_handed_perspective( - vertical_field_of_view: np.ndarray, - aspect_ratio: np.ndarray, - near: np.ndarray, - far: np.ndarray, -) -> np.ndarray: + vertical_field_of_view: FloatArray, + aspect_ratio: FloatArray, + near: FloatArray, + far: FloatArray, +) -> FloatArray: """Builds a right handed perspective projection matrix. Similar to tensorflow_graphics.rendering.camera.perspective.right_handed. @@ -759,7 +761,6 @@ def _right_handed_perspective( Returns: The perspective projection matrix, (..., 4, 4) in OpenGL convention. """ - itan_half_vertical_field_of_view = 1.0 / np.tan(vertical_field_of_view * 0.5) zero = np.zeros_like(itan_half_vertical_field_of_view) one = np.ones_like(itan_half_vertical_field_of_view) @@ -807,10 +808,10 @@ def load_edgeflow_texture(texture_path: epath.Path | None) -> FloatArray | None: def _vertex_group_mean( - vertices: np.ndarray, + vertices: FloatArray, group_names: Sequence[str], gnm_np: gnm_numpy.GNM, -): +) -> FloatArray: """Gets the average point of GNM vertex groups. Args: @@ -853,7 +854,7 @@ def load_texture( if edgeflow_path is not None: edgeflow = load_edgeflow_texture(edgeflow_path) if edgeflow is not None: - texture_dict[SKIN_UDIM_TILE] = edgeflow[..., None] + texture_dict[SKIN_UDIM_TILE] = edgeflow[..., np.newaxis] elif isinstance(texture, np.ndarray): texture_dict[SKIN_UDIM_TILE] = texture elif isinstance(texture, Mapping): @@ -867,11 +868,12 @@ def load_texture( f' {gnm_tiles}.' ) - def _to_3channel_uint8(texture_image: np.ndarray) -> npt.NDArray[np.uint8]: + def _to_3channel_uint8( + texture_image: FloatArray | npt.NDArray[np.uint8], + ) -> npt.NDArray[np.uint8]: if texture_image.dtype != np.uint8: - texture_image = (np.clip(texture_image, 0.0, 1.0) * 255.0).astype( - np.uint8 - ) + texture_image = np.clip(texture_image, 0.0, 1.0) * 255.0 + texture_image = texture_image.astype(np.uint8, copy=False) if texture_image.shape[-1] == 1: texture_image = np.repeat(texture_image, 3, axis=-1) return texture_image @@ -901,9 +903,23 @@ def _udim_tiles_of_texture_key(key: int | np.integer | str) -> list[int]: return tiles +@overload +def _adjust_scalar_shape( + array: npt.NDArray[np.bool_] | bool, batch_dims: Sequence[int] +) -> npt.NDArray[np.bool_]: + ... + + +@overload def _adjust_scalar_shape( - array: np.ndarray | float, batch_dims: Sequence[int] -) -> np.ndarray: + array: FloatArray | float, batch_dims: Sequence[int] +) -> FloatArray: + ... + + +def _adjust_scalar_shape( + array: npt.NDArray[np.generic] | float, batch_dims: Sequence[int] +) -> npt.NDArray[np.generic]: """Adjusts the shape of a scalar to match the batch dimensions. Args: @@ -913,16 +929,17 @@ def _adjust_scalar_shape( Returns: The adjusted array, (..., 1). """ - # If the array is a scalar or 1D array, add a final dimension of size 1. array = np.array(array) if array.ndim <= 1: - array = array[..., None] + array = array[..., np.newaxis] return np.broadcast_to(array, (*batch_dims, 1)) -def get_batch_dim(*arrays: tuple[np.ndarray | None, int]) -> tuple[int, ...]: +def get_batch_dim( + *arrays: tuple[npt.NDArray[np.generic] | None, int], +) -> tuple[int, ...]: """Finds the largest batch dimension that all arrays can be broadcast to. Requires that all arrays have batch dimensions that can be broadcast together. @@ -936,6 +953,7 @@ def get_batch_dim(*arrays: tuple[np.ndarray | None, int]) -> tuple[int, ...]: Returns: The batch dimensions that all arrays can be safely broadcast to. + Raises: ValueError: If the arrays cannot be broadcast together. """ diff --git a/gnm/shape/visualization/render_common_test.py b/gnm/shape/visualization/render_common_test.py index ec73f3a7..ebe2442e 100644 --- a/gnm/shape/visualization/render_common_test.py +++ b/gnm/shape/visualization/render_common_test.py @@ -26,6 +26,7 @@ from gnm.shape.visualization import camera_conversions from gnm.shape.visualization import render_common import numpy as np +import numpy.typing as npt _TupleOfInts = tuple[int, ...] @@ -44,7 +45,7 @@ def _load_head_gnms() -> dict[str, gnm_numpy.GNM]: def _get_random_parameters( batch_dims: _TupleOfInts, gnm_np: gnm_numpy.GNM, -) -> dict[str, np.ndarray]: +) -> dict[str, npt.NDArray[np.float32]]: """Returns random GNM parameters. Args: @@ -64,12 +65,12 @@ def _get_random_parameters( -0.2, 0.2, size=batch_dims + (gnm_np.num_joints, 3) ) translation = np.random.uniform(-0.5, 0.5, size=batch_dims + (3,)) * 0.0 - return dict( - identity=identity.astype(np.float32), - expression=expression.astype(np.float32), - rotations=rotations.astype(np.float32), - translation=translation.astype(np.float32), - ) + return { + 'identity': identity.astype(np.float32), + 'expression': expression.astype(np.float32), + 'rotations': rotations.astype(np.float32), + 'translation': translation.astype(np.float32), + } class TestGetBatchDim(parameterized.TestCase): @@ -443,7 +444,7 @@ def mock_backend(*args: Any, **kwargs: Any) -> render_common.FloatArray: w, h = kwargs['image_size'] return np.zeros((n, h, w, 3), dtype=np.float32) - w2c_opencv = np.eye(4, dtype=np.float32)[None, ...] + w2c_opencv = np.eye(4, dtype=np.float32)[np.newaxis] # Without conversion render_common.render_gnm_mesh(