From bd5d087f91e4ff4b6e3e9e5d147bc1de6a164daf Mon Sep 17 00:00:00 2001 From: Mirco Mazzoni Date: Fri, 2 Oct 2026 03:30:29 -0700 Subject: [PATCH] Dedupe test model lookups. PiperOrigin-RevId: 992220138 --- .../fitting_utils/project_on_pca_test.py | 53 ++---- gnm/shape/gnm_jax_test.py | 77 +++----- gnm/shape/gnm_numpy_test.py | 148 +++++----------- gnm/shape/gnm_pytorch_test.py | 93 +++------- gnm/shape/gnm_tensorflow_test.py | 95 ++++------ gnm/shape/gnm_utils_test.py | 164 ++++++------------ gnm/shape/visualization/gnm_mitsuba_test.py | 6 +- gnm/shape/visualization/gnm_pyrender_test.py | 6 +- gnm/shape/visualization/render_common_test.py | 58 +++---- .../visualization/render_gnm_parity_test.py | 8 +- gnm/shape/visualization/render_gnm_test.py | 10 +- .../visualization/render_gnm_test_base.py | 36 ++-- gnm/shape/visualization/vertex_colors_test.py | 36 ++-- 13 files changed, 264 insertions(+), 526 deletions(-) diff --git a/gnm/shape/fitting_utils/project_on_pca_test.py b/gnm/shape/fitting_utils/project_on_pca_test.py index 7eb04e90..60b735a2 100644 --- a/gnm/shape/fitting_utils/project_on_pca_test.py +++ b/gnm/shape/fitting_utils/project_on_pca_test.py @@ -19,6 +19,7 @@ from absl.testing import absltest from absl.testing import parameterized from gnm.shape import gnm_numpy +from gnm.shape import gnm_test_utils from gnm.shape import gnm_utils from gnm.shape.data.versions import gnm_test_catalog from gnm.shape.fitting_utils import project_on_pca @@ -37,16 +38,20 @@ def setUpClass(cls): super().setUpClass() cls.rng = np.random.default_rng(0) - cls.gnms = {} - for version in gnm_test_catalog.MAINTAINED_MAJOR_VERSIONS: - if ( - gnm_numpy.GNMVariant.HEAD - in gnm_test_catalog.MAJOR_VERSION_TO_VARIANTS_MAP[version] - ): - cls.gnms[version] = gnm_numpy.GNM.from_remote( - gnm_numpy.GNMMajorVersion(version.removeprefix('v')), - gnm_numpy.GNMVariant.HEAD, + cls.gnms = { + version: gnm_test_utils.load_gnm( + gnm_numpy.GNM, version, gnm_numpy.GNMVariant.HEAD ) + for version in gnm_test_catalog.MAINTAINED_MAJOR_VERSIONS + if gnm_numpy.GNMVariant.HEAD + in gnm_test_catalog.MAJOR_VERSION_TO_VARIANTS_MAP[version] + } + + def _get_gnm(self, version: str) -> gnm_numpy.GNM: + """Returns the head model of `version`, or skips if it is unavailable.""" + if version not in self.gnms: + self.skipTest(f'Head variant not available in version {version}') + return self.gnms[version] @parameterized.product( version=gnm_test_catalog.MAINTAINED_MAJOR_VERSIONS, @@ -55,7 +60,7 @@ def setUpClass(cls): ) def test_can_recover_identity(self, version: str, index: int, value: float): """Tests that the estimated identity is close to the ground truth.""" - gnm_np = self.gnms[version] + gnm_np = self._get_gnm(version) identity = np.zeros(gnm_np.identity_dim) expression = np.zeros(gnm_np.expression_dim) identity[index] = value @@ -78,7 +83,7 @@ def test_can_recover_identity(self, version: str, index: int, value: float): ) def test_can_recover_expression(self, version: str, index: int, value: float): """Tests that the estimated expression is close to the ground truth.""" - gnm_np = self.gnms[version] + gnm_np = self._get_gnm(version) identity = np.zeros(gnm_np.identity_dim) expression = np.zeros(gnm_np.expression_dim) expression[index] = value @@ -102,17 +107,10 @@ def test_can_fit_identity_using_subset_of_basis( self, version, num_components_and_threshold ): """Tests that we can estimate identity using a subset of the basis.""" - if ( - gnm_numpy.GNMVariant.HEAD - not in gnm_test_catalog.MAJOR_VERSION_TO_VARIANTS_MAP[version] - ): - self.skipTest(f'Head variant not available in version {version}') + gnm = self._get_gnm(version) num_components, threshold = num_components_and_threshold - major_version = gnm_numpy.GNMMajorVersion(version.removeprefix('v')) - gnm = gnm_numpy.GNM.from_remote(major_version, gnm_numpy.GNMVariant.HEAD) - identity = self.rng.uniform(low=-1.5, high=1.5, size=gnm.identity_dim) expression = np.zeros(gnm.expression_dim) @@ -149,17 +147,10 @@ def test_can_recover_identity_with_object( self, version, batch_size_num_components_threshold ): """Tests that the projection object can estimate identity .""" - if ( - gnm_numpy.GNMVariant.HEAD - not in gnm_test_catalog.MAJOR_VERSION_TO_VARIANTS_MAP[version] - ): - self.skipTest(f'Head variant not available in version {version}') + gnm = self._get_gnm(version) batch_size, num_components, threshold = batch_size_num_components_threshold - major_version = gnm_numpy.GNMMajorVersion(version.removeprefix('v')) - gnm = gnm_numpy.GNM.from_remote(major_version, gnm_numpy.GNMVariant.HEAD) - identity = self.rng.uniform( low=-1.5, high=1.5, size=(batch_size, gnm.identity_dim) ) @@ -204,15 +195,9 @@ def test_can_recover_expression_with_object( self, version, batch_size_num_components_threshold ): """Tests that the projection object can estimate expression.""" - if ( - gnm_numpy.GNMVariant.HEAD - not in gnm_test_catalog.MAJOR_VERSION_TO_VARIANTS_MAP[version] - ): - self.skipTest(f'Head variant not available in version {version}') + gnm = self._get_gnm(version) batch_size, num_components, threshold = batch_size_num_components_threshold - major_version = gnm_numpy.GNMMajorVersion(version.removeprefix('v')) - gnm = gnm_numpy.GNM.from_remote(major_version, gnm_numpy.GNMVariant.HEAD) expected_regions = [ 'left_eye_region', diff --git a/gnm/shape/gnm_jax_test.py b/gnm/shape/gnm_jax_test.py index 6625370d..9e7cfea1 100644 --- a/gnm/shape/gnm_jax_test.py +++ b/gnm/shape/gnm_jax_test.py @@ -81,6 +81,15 @@ def setUpClass(cls): ) cls.addClassCleanup(cls.gnms_jax.clear) + def _get_gnms( + self, version: str, variant: Any + ) -> tuple[gnm_jax.GNM, gnm_numpy.GNM]: + """Returns the JAX and NumPy GNMs, or skips the test if unavailable.""" + return ( + self.gnms_jax.get_or_skip(version, variant), + self.gnms_np.get_or_skip(version, variant), + ) + def _get_default_kwargs( self, gnm_np: gnm_numpy.GNM, n_batch: int = 1 ) -> dict[str, jnp.ndarray]: @@ -118,11 +127,7 @@ def _get_random_kwargs( ) def test_properties_match(self, version: str, variant: Any): """Tests that important properties match between implementations.""" - variant_str = variant.value if hasattr(variant, 'value') else variant - if variant_str not in self.gnms_np[version]: - self.skipTest(f'variant {variant_str} not supported in {version}.') - gnm = self.gnms_jax[version][variant_str] - gnm_np = self.gnms_np[version][variant_str] + gnm, gnm_np = self._get_gnms(version, variant) self.assertEqual(gnm.version, gnm_np.version) self.assertEqual(gnm.identity_dim, gnm_np.identity_dim) @@ -148,11 +153,7 @@ def test_parity_with_gnm_numpy( dtype: jnp.dtype, ): """Tests that JAX GNM poses vertices the same as NumPy GNM.""" - variant_str = variant.value if hasattr(variant, 'value') else variant - if variant_str not in self.gnms_np[version]: - self.skipTest(f'variant {variant_str} not supported in {version}.') - gnm = self.gnms_jax[version][variant_str] - gnm_np = self.gnms_np[version][variant_str] + gnm, gnm_np = self._get_gnms(version, variant) # Build a batch of random parameters. gnm_jax_parameters = self._get_random_kwargs(gnm_np, batch_dim) @@ -177,11 +178,7 @@ def test_vertices_and_landmarks( batch_size: tuple[int, ...], ): """Tests extracting vertices and landmarks in JAX.""" - variant_str = variant.value if hasattr(variant, 'value') else variant - if variant_str not in self.gnms_np[version]: - self.skipTest(f'variant {variant_str} not supported in {version}.') - gnm = self.gnms_jax[version][variant_str] - gnm_np = self.gnms_np[version][variant_str] + gnm, gnm_np = self._get_gnms(version, variant) gnm_jax_parameters = self._get_random_kwargs(gnm_np, batch_size) verts, landmarks = gnm.vertices_and_landmarks( @@ -197,11 +194,7 @@ def test_vertices_and_landmarks( dtype=DTYPES, ) def test_jax_jit_behavior(self, version: str, variant: Any, dtype: jnp.dtype): - variant_str = variant.value if hasattr(variant, 'value') else variant - if variant_str not in self.gnms_np[version]: - self.skipTest(f'variant {variant_str} not supported in {version}.') - gnm = self.gnms_jax[version][variant_str] - gnm_np = self.gnms_np[version][variant_str] + gnm, gnm_np = self._get_gnms(version, variant) # Build a batch of random parameters. n_batch = 10 @@ -234,11 +227,7 @@ def jax_jit_wrapper(identity, expression, rotations, translation): ) def test_bad_shape(self, version: str, variant: Any): """Badly shaped parameter should throw an error.""" - variant_str = variant.value if hasattr(variant, 'value') else variant - if variant_str not in self.gnms_np[version]: - self.skipTest(f'variant {variant_str} not supported in {version}.') - gnm = self.gnms_jax[version][variant_str] - gnm_np = self.gnms_np[version][variant_str] + gnm, gnm_np = self._get_gnms(version, variant) bad_dimension = gnm.expression_dim + gnm.identity_dim + gnm.num_joints n_batch = 5 @@ -254,11 +243,7 @@ def test_bad_shape(self, version: str, variant: Any): ) def test_bad_shape_joint_transforms(self, version: str, variant: Any): """Badly shaped parameter should throw an error.""" - variant_str = variant.value if hasattr(variant, 'value') else variant - if variant_str not in self.gnms_np[version]: - self.skipTest(f'variant {variant_str} not supported in {version}.') - gnm = self.gnms_jax[version][variant_str] - gnm_np = self.gnms_np[version][variant_str] + gnm, gnm_np = self._get_gnms(version, variant) bad_dimension = gnm.identity_dim + gnm.num_joints bad_inputs = [ @@ -284,11 +269,7 @@ def test_joint_transforms_numpy_parity( self, version: str, variant: Any, dtype: jnp.dtype ): """Test that the joint transformations function matches Numpy.""" - variant_str = variant.value if hasattr(variant, 'value') else variant - if variant_str not in self.gnms_np[version]: - self.skipTest(f'variant {variant_str} not supported in {version}.') - gnm = self.gnms_jax[version][variant_str] - gnm_np = self.gnms_np[version][variant_str] + gnm, gnm_np = self._get_gnms(version, variant) # Build a batch of random parameters. n_batch = 10 @@ -330,11 +311,7 @@ def test_gradients( # pyrefly: ignore[bad-function-definition] dtype: jnp.dtype = jnp.float32, ): - variant_str = variant.value if hasattr(variant, 'value') else variant - if variant_str not in self.gnms_np[version]: - self.skipTest(f'variant {variant_str} not supported in {version}.') - gnm = self.gnms_jax[version][variant_str] - gnm_np = self.gnms_np[version][variant_str] + gnm, gnm_np = self._get_gnms(version, variant) # Build a batch of random parameters. grad_func = jax.grad( @@ -388,11 +365,7 @@ def test_gradients( ) def test_transformed_local_offsets(self, version: str, variant: Any): """Tests that vertex offset is applied to the posed positions.""" - variant_str = variant.value if hasattr(variant, 'value') else variant - if variant_str not in self.gnms_np[version]: - self.skipTest(f'variant {variant_str} not supported in {version}.') - gnm_np = self.gnms_np[version][variant_str] - gnm = self.gnms_jax[version][variant_str] + gnm, gnm_np = self._get_gnms(version, variant) gnm_jax_parameters = self._get_default_kwargs(n_batch=10, gnm_np=gnm_np) local_vertices = gnm.vertex_positions_bind_pose( @@ -433,10 +406,7 @@ def test_transformed_local_offsets(self, version: str, variant: Any): ) def test_to_numpy_data_dict(self, version: str, variant: Any): """Tests to_numpy_data_dict method.""" - variant_str = variant.value if hasattr(variant, 'value') else variant - if variant_str not in self.gnms_np[version]: - self.skipTest(f'variant {variant_str} not supported in {version}.') - gnm = self.gnms_jax[version][variant_str] + gnm = self.gnms_jax.get_or_skip(version, variant) data_dict = gnm.to_numpy_data_dict() @@ -459,19 +429,14 @@ def test_to_numpy_data_dict(self, version: str, variant: Any): ) def test_from_gnm(self, version: str, variant: Any): """Tests from_gnm factory method.""" - variant_str = variant.value if hasattr(variant, 'value') else variant - if variant_str not in self.gnms_np[version]: - self.skipTest(f'variant {variant_str} not supported in {version}.') - gnm = self.gnms_jax[version][variant_str] + gnm, gnm_np = self._get_gnms(version, variant) new_gnm = gnm_jax.GNM.from_gnm(gnm) self.assertEqual(new_gnm.version, gnm.version) self.assertEqual(new_gnm.variant, gnm.variant) - parameters = self._get_random_kwargs( - self.gnms_np[version][variant_str], n_batch=(10,) - ) + parameters = self._get_random_kwargs(gnm_np, n_batch=(10,)) parameters_jax = jax.tree.map(jnp.asarray, parameters) vertices_orig = gnm(**parameters_jax) diff --git a/gnm/shape/gnm_numpy_test.py b/gnm/shape/gnm_numpy_test.py index 2d4311f1..2be009d6 100644 --- a/gnm/shape/gnm_numpy_test.py +++ b/gnm/shape/gnm_numpy_test.py @@ -178,10 +178,7 @@ def _get_default_kwargs( ) def test_global_translation(self, version: str, variant: str): """Test that global translation shifts joints and vertices uniformly.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - - gnm_np = self.gnms[version][variant] + gnm_np = self.gnms.get_or_skip(version, variant) translation = np.array([1, 2, 3], dtype=np.float32) kwargs = self._get_default_kwargs(gnm_np) desired = gnm_np(**kwargs) + translation @@ -209,9 +206,7 @@ def test_global_translation(self, version: str, variant: str): ) def test_global_rotation(self, version: str, variant: str): """Test that a root joint rotation is correctly modelled.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms[version][variant] + gnm_np = self.gnms.get_or_skip(version, variant) kwargs = self._get_default_kwargs(gnm_np) vertices_zero_pose = gnm_np(**kwargs) @@ -263,9 +258,7 @@ def test_global_rotation(self, version: str, variant: str): ) def test_bad_shape(self, version: str, variant: str): """Badly shaped parameter should throw a ValueError.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms[version][variant] + gnm_np = self.gnms.get_or_skip(version, variant) bad_dimension = ( gnm_np.expression_dim + gnm_np.identity_dim + gnm_np.num_joints ) @@ -280,9 +273,7 @@ def test_bad_shape(self, version: str, variant: str): variant=tuple(_SUPPORTED_VARIANTS), ) def test_bad_shape_joint_transforms(self, version: str, variant: str): - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms[version][variant] + gnm_np = self.gnms.get_or_skip(version, variant) bad_dimension = ( gnm_np.expression_dim + gnm_np.identity_dim + gnm_np.num_joints ) @@ -307,9 +298,7 @@ def test_bad_shape_joint_transforms(self, version: str, variant: str): ) def test_identity(self, version: str, variant: str): """Test that vertices are different after applying the identity model.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms[version][variant] + gnm_np = self.gnms.get_or_skip(version, variant) kwargs = self._get_default_kwargs(gnm_np) verts_zero_identity = gnm_np(**kwargs) identity = np.ones(gnm_np.identity_dim) @@ -324,9 +313,7 @@ def test_identity(self, version: str, variant: str): ) def test_expression(self, version: str, variant: str): """Test that vertices are different after applying the expression model.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms[version][variant] + gnm_np = self.gnms.get_or_skip(version, variant) kwargs = self._get_default_kwargs(gnm_np) verts_zero_expression = gnm_np(**kwargs) expression = np.ones(gnm_np.expression_dim) @@ -342,9 +329,7 @@ def test_vertices_and_landmarks( self, version: str, variant: str, batch_size: tuple[int, ...] ): """Test extracting vertices and landmarks.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms[version][variant] + gnm_np = self.gnms.get_or_skip(version, variant) kwargs = self._get_default_kwargs(gnm_np, batch_dims=batch_size) verts, landmarks = gnm_np.vertices_and_landmarks( gnm_numpy.GNMLandmarksType.HEAD_SPARSE_68, **kwargs @@ -358,9 +343,8 @@ def test_vertices_and_landmarks( ) def test_vertex_groups_exist(self, version: str, variant: str): """Tests that there is at least one vertex group defined.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - self.assertNotEmpty(self.gnms[version][variant].vertex_group_names) + gnm_np = self.gnms.get_or_skip(version, variant) + self.assertNotEmpty(gnm_np.vertex_group_names) @parameterized.product( version=_MAINTAINED_MAJOR_GNM_VERSIONS, @@ -368,9 +352,7 @@ def test_vertex_groups_exist(self, version: str, variant: str): ) def test_vertex_groups(self, version: str, variant: str): """Tests that we can retrieve values and indices for each vertex group.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms[version][variant] + gnm_np = self.gnms.get_or_skip(version, variant) # Skip the groups that are zero by design, e.g. the finger tips. valid_group_names = [ @@ -397,9 +379,7 @@ def test_vertex_groups(self, version: str, variant: str): ) def test_uvs(self, version: str, variant: str): """Tests that texture coordinates have the correct format.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) self.assertEqual(gnm.quad_uvs.shape[:2], gnm.quads.shape) self.assertEqual(gnm.triangle_uvs.shape[:2], gnm.triangles.shape) self.assertEqual(gnm.vertex_uvs.shape, (gnm.num_vertices, 2)) @@ -424,9 +404,7 @@ def test_uvs(self, version: str, variant: str): ) def test_edge_list(self, version: str, variant: str): """Checks that the edge list matches the quad topology.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms[version][variant] + gnm_np = self.gnms.get_or_skip(version, variant) # Every quad side must appear in the edge list in both directions, and the # edge list must not contain any other edge. Compare the sets of unique @@ -446,9 +424,7 @@ def test_edge_list(self, version: str, variant: str): @parameterized.parameters(get_group_subsets_test_cases()) def test_group_subsets(self, version: str, variant: str, group_name: str): """Tests that the convenience functions for accessing data subsets work.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms[version][variant] + gnm_np = self.gnms.get_or_skip(version, variant) group_indices = gnm_np.vertex_group_indices(group_name) triangles = gnm_np.triangles_group(group_name) @@ -480,9 +456,7 @@ def test_parts_are_non_overlapping_and_complete( self, version: str, variant: str ): """Check that separate parts are non-overlapping and complete GNM.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) combined = np.stack([gnm.vertex_group(n) for n in gnm.mesh_component_names]) np.testing.assert_array_equal( np.sum(combined, axis=0), np.ones(gnm.num_vertices) @@ -494,9 +468,7 @@ def test_parts_are_non_overlapping_and_complete( ) def test_mesh_component_names_include_skin(self, version: str, variant: str): """Check that the mesh_component_names include 'skin'.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) self.assertIn('skin', gnm.mesh_component_names) @parameterized.product( @@ -505,9 +477,7 @@ def test_mesh_component_names_include_skin(self, version: str, variant: str): ) def test_vertex_group_names_include_skin(self, version: str, variant: str): """Check that the vertex_group_names include 'skin'.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) self.assertIn('skin', gnm.vertex_group_names) @parameterized.product( @@ -521,9 +491,7 @@ def test_mirror_indices(self, version: str, variant: str): See mirror_indices_test. """ - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) self.assertCountEqual(gnm.mirror_indices, range(gnm.num_vertices)) @parameterized.product( @@ -535,9 +503,7 @@ def test_mirror_indices(self, version: str, variant: str): ) def test_eyes_dont_move(self, version: str, variant: str, vertex_group: str): """Eyes should not move even if every expression shape is applied.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms[version][variant] + gnm_np = self.gnms.get_or_skip(version, variant) indices = gnm_np.vertex_group_indices(vertex_group) all_expressions_no_pupil = np.ones(gnm_np.expression_dim) regions = gnm_utils.expression_to_regions(all_expressions_no_pupil, gnm_np) @@ -555,9 +521,7 @@ def test_joints_with_identity_dont_match_template( self, version: str, variant: str ): """Given a random identity, joints should _not_ match the template.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms[version][variant] + gnm_np = self.gnms.get_or_skip(version, variant) thresholds = get_joints_with_identity_thresholds(gnm_np) @@ -577,9 +541,7 @@ def test_eye_interiors_is_inside_exterior( self, version: str, variant: str, side: str ): """Tests eye interior is inside exterior for template and identities.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms[version][variant] + gnm_np = self.gnms.get_or_skip(version, variant) interior_indices = gnm_np.vertex_group_indices('eye_interiors', side) exterior_triangles = gnm_np.triangles_group('eye_exteriors', side) @@ -615,9 +577,7 @@ def test_eye_interiors_is_inside_exterior( ) def test_expression_basis_shape(self, version: str, variant: str): """Tests the expression basis is accessible and has the expected shape.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms[version][variant] + gnm_np = self.gnms.get_or_skip(version, variant) self.assertEqual( gnm_np.expression_basis.shape, (gnm_np.expression_dim, gnm_np.num_vertices, 3), @@ -637,9 +597,7 @@ def test_pupil_uvs_close_to_middle( self, version: str, variant: str, vertex_groups: Sequence[str] ): """Tests the pupil UV coordinates are close to (0.5, 0.5).""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms[version][variant] + gnm_np = self.gnms.get_or_skip(version, variant) uv = np.mean(gnm_np.vertex_uvs_group(*vertex_groups), axis=0) self.assertLess(np.linalg.norm(uv - 0.5), 2e-3) @@ -656,9 +614,7 @@ def test_eye_uvs_not_flipped( self, version: str, variant: str, side: str, group: str, axis: int ): """Checks that model-space axes align with UV axes for eyes.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms[version][variant] + gnm_np = self.gnms.get_or_skip(version, variant) side_indices = gnm_np.vertex_group_indices(side) group_indices = gnm_np.vertex_group_indices(group) side_group_indices = np.intersect1d(side_indices, group_indices) @@ -675,9 +631,7 @@ def test_eye_uvs_not_flipped( ) def test_omit_all_parameters(self, version: str, variant: str): """If all parameters are omitted, GNM returns the template vertices.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) expected = gnm.template_vertex_positions np.testing.assert_allclose(expected, gnm(), atol=1e-6) @@ -695,9 +649,7 @@ def test_variable_batch_and_omitted_parameters( parameter_count: int, ): """Exercise GNM with various batch dimensions and omitted parameters.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) parameters = self._get_default_kwargs(gnm, batch_dims) @@ -712,9 +664,7 @@ def test_variable_batch_and_omitted_parameters( ) def test_pose_correctives(self, version: str, variant: str): """Test that the pose correctives are correct.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms[version][variant] + gnm_np = self.gnms.get_or_skip(version, variant) test_cases = get_pose_correctives_test_cases(gnm_np) for case in test_cases: @@ -733,9 +683,7 @@ def test_pose_correctives(self, version: str, variant: str): ) def test_multiple_vertex_groups(self, version: str, variant: str): """Test we can combine multiple vertex groups.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) # Compare two ways of combining both eyeballs fully. np.testing.assert_array_equal( @@ -776,9 +724,7 @@ def test_left_and_right_have_equal_number_of_vertices( self, version: str, variant: str ): """Test the left and right groups have the same number of vertices.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) self.assertEqual( len(gnm.vertex_group_indices('left')), len(gnm.vertex_group_indices('right')), @@ -792,9 +738,7 @@ def test_face_identity_names_contain_all_groups( self, version: str, variant: str ): """Tests that the identity names contain all expected groups.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) # Only check for the groups existing in the GNM model. expected_groups = gnm_utils.identity_region_names(gnm) @@ -828,9 +772,7 @@ def test_face_expression_names_contain_all_groups( self, version: str, variant: str ): """Tests that the expression names contain all expected groups.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) regexes = [ f'^{n}_[0-9][0-9][0-9]$' @@ -902,9 +844,8 @@ def test_adding_and_removing_vertex_group_names_does_not_raise_error( self, version: str, variant: str ): """Tests that adding and removing vertex group names does not raise.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + # Copy to avoid modifying the cached model for other tests. + gnm = copy.copy(self.gnms.get_or_skip(version, variant)) gnm.vertex_group_names = [*gnm.vertex_group_names, 'new_group'] gnm.vertex_group_names = gnm.vertex_group_names[:-1] @@ -914,9 +855,7 @@ def test_adding_and_removing_vertex_group_names_does_not_raise_error( ) def test_add_vertex_group(self, version: str, variant: str): """Tests adding a new vertex group.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) new_group_name = 'test_new_group' new_group_value = np.random.uniform(size=(gnm.num_vertices,)) @@ -945,9 +884,7 @@ def test_add_vertex_group(self, version: str, variant: str): ) def test_to_numpy_data_dict(self, version: str, variant: str): """Tests to_numpy_data_dict method.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) data_dict = gnm.to_numpy_data_dict() @@ -970,9 +907,7 @@ def test_to_numpy_data_dict(self, version: str, variant: str): ) def test_from_gnm(self, version: str, variant: str): """Tests from_gnm factory method.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) new_gnm = gnm_numpy.GNM.from_gnm(gnm) @@ -993,9 +928,7 @@ def test_compute_vertex_normals( self, version: str, variant: str, batch_dims: Sequence[int] ): """Tests that the vertex normals are computed correctly.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms[version][variant] + gnm_np = self.gnms.get_or_skip(version, variant) parameters = self._get_default_kwargs(gnm_np, batch_dims=batch_dims) vertices = gnm_np(**parameters) vertex_normals = gnm_np.compute_vertex_normals(vertices) @@ -1026,10 +959,9 @@ def test_compute_vertex_normals_with_zero_magnitude( self, version: str, variant: str ): """Tests that a warning is logged for vertex normals with zero magnitude.""" - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms[version][variant] - vertices = gnm_np.template_vertex_positions + gnm_np = self.gnms.get_or_skip(version, variant) + # Copy, so that the cached model shared with other tests isn't modified. + vertices = gnm_np.template_vertex_positions.copy() # Collapse the faces around vertex 0. tmesh = trimesh.Trimesh(vertices, gnm_np.triangles, process=False) @@ -1052,7 +984,7 @@ def test_benchmark_call( max_runtime_ms: float, ): """Tests the benchmark for GNM on Float32 operations.""" - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) parameters = self._get_default_kwargs(gnm, batch_dims=[batch_size]) parameters = {k: v.astype(np.float32) for k, v in parameters.items()} diff --git a/gnm/shape/gnm_pytorch_test.py b/gnm/shape/gnm_pytorch_test.py index 506a1b59..ea14dc0c 100644 --- a/gnm/shape/gnm_pytorch_test.py +++ b/gnm/shape/gnm_pytorch_test.py @@ -15,6 +15,7 @@ """Tests for GNM PyTorch implementation.""" from collections.abc import Sequence +import copy import functools import itertools from typing import Any @@ -78,6 +79,15 @@ def setUpClass(cls): ) cls.addClassCleanup(cls.gnms_pytorch.clear) + def _get_gnms( + self, version: str, variant: Any + ) -> tuple[gnm_pytorch.GNM, gnm_numpy.GNM]: + """Returns the PyTorch and NumPy GNMs, or skips the test if unavailable.""" + return ( + self.gnms_pytorch.get_or_skip(version, variant), + self.gnms_np.get_or_skip(version, variant), + ) + def _get_default_kwargs( self, gnm_np: gnm_numpy.GNM, n_batch: int = 1, device: str = 'cpu' ) -> dict[str, torch.Tensor]: @@ -123,11 +133,7 @@ def _get_random_kwargs( ) def test_properties_match(self, version: str, variant: Any): """Tests that important properties match between implementations.""" - variant_str = variant.value if hasattr(variant, 'value') else variant - if variant_str not in self.gnms_np[version]: - self.skipTest(f'variant {variant_str} not supported in {version}.') - gnm_torch = self.gnms_pytorch[version][variant_str] - gnm_np = self.gnms_np[version][variant_str] + gnm_torch, gnm_np = self._get_gnms(version, variant) self.assertEqual(gnm_torch.version, gnm_np.version) self.assertEqual(gnm_torch.identity_dim, gnm_np.identity_dim) @@ -141,11 +147,7 @@ def test_properties_match(self, version: str, variant: Any): ) def test_parity_with_gnm_numpy(self, version: str, variant: Any): """Tests that PyTorch GNM poses vertices the same as NumPy GNM.""" - variant_str = variant.value if hasattr(variant, 'value') else variant - if variant_str not in self.gnms_np[version]: - self.skipTest(f'variant {variant_str} not supported in {version}.') - gnm_np = self.gnms_np[version][variant_str] - gnm_torch = self.gnms_pytorch[version][variant_str] + gnm_torch, gnm_np = self._get_gnms(version, variant) # Build a batch of random parameters. n_batch = 10 @@ -172,11 +174,7 @@ def test_vertices_and_landmarks( self, version: str, variant: Any, batch_size: tuple[int, ...] ): """Tests extracting vertices and landmarks in PyTorch.""" - variant_str = variant.value if hasattr(variant, 'value') else variant - if variant_str not in self.gnms_np[version]: - self.skipTest(f'variant {variant_str} not supported in {version}.') - gnm_torch = self.gnms_pytorch[version][variant_str] - gnm_np = self.gnms_np[version][variant_str] + gnm_torch, gnm_np = self._get_gnms(version, variant) parameters_torch = self._get_random_kwargs( gnm_np, batch_size, device=self.device @@ -193,11 +191,7 @@ def test_vertices_and_landmarks( ) def test_bad_shape(self, version: str, variant: Any): """Badly shaped parameter should throw an error.""" - variant_str = variant.value if hasattr(variant, 'value') else variant - if variant_str not in self.gnms_np[version]: - self.skipTest(f'variant {variant_str} not supported in {version}.') - gnm_np = self.gnms_np[version][variant_str] - gnm_torch = self.gnms_pytorch[version][variant_str] + gnm_torch, gnm_np = self._get_gnms(version, variant) bad_dimension = ( gnm_torch.expression_dim + gnm_torch.identity_dim + gnm_np.num_joints @@ -215,11 +209,7 @@ def test_bad_shape(self, version: str, variant: Any): ) def test_bad_shape_joint_transforms(self, version: str, variant: Any): """Badly shaped parameter should throw an error.""" - variant_str = variant.value if hasattr(variant, 'value') else variant - if variant_str not in self.gnms_np[version]: - self.skipTest(f'variant {variant_str} not supported in {version}.') - gnm_np = self.gnms_np[version][variant_str] - gnm_torch = self.gnms_pytorch[version][variant_str] + gnm_torch, gnm_np = self._get_gnms(version, variant) bad_dimension = ( gnm_torch.expression_dim + gnm_torch.identity_dim + gnm_np.num_joints @@ -248,11 +238,7 @@ def test_bad_shape_joint_transforms(self, version: str, variant: Any): ) def test_joint_transforms_numpy_parity(self, version: str, variant: Any): """Test that the joint transformations function matches NumPy.""" - variant_str = variant.value if hasattr(variant, 'value') else variant - if variant_str not in self.gnms_np[version]: - self.skipTest(f'variant {variant_str} not supported in {version}.') - gnm_np = self.gnms_np[version][variant_str] - gnm_torch = self.gnms_pytorch[version][variant_str] + gnm_torch, gnm_np = self._get_gnms(version, variant) # Build a batch of random parameters. n_batch = 10 @@ -283,13 +269,10 @@ def test_joint_transforms_numpy_parity(self, version: str, variant: Any): variant=tuple(_SUPPORTED_VARIANTS), ) def test_prune_vertices(self, version: str, variant: Any): - variant_str = variant.value if hasattr(variant, 'value') else variant - if variant_str not in self.gnms_np[version]: - self.skipTest(f'variant {variant_str} not supported in {version}.') - gnm_np = self.gnms_np[version][variant_str] - gnm_torch = self.gnms_pytorch[version][variant_str] + gnm_torch, gnm_np = self._get_gnms(version, variant) - gnm_pruned = gnm_test_utils.load_gnm(gnm_pytorch.GNM, version, variant_str) + # Prune a copy to avoid modifying the cached model for other tests. + gnm_pruned = copy.deepcopy(gnm_torch) keep_vertices = gnm_np.quads[0] gnm_pruned.prune_vertices(keep_vertices) @@ -315,11 +298,7 @@ def test_prune_vertices(self, version: str, variant: Any): ) def test_no_batch(self, version: str, variant: Any): """Tests we can use Torch GNM without a leading batch dimension.""" - variant_str = variant.value if hasattr(variant, 'value') else variant - if variant_str not in self.gnms_np[version]: - self.skipTest(f'variant {variant_str} not supported in {version}.') - gnm_np = self.gnms_np[version][variant_str] - gnm_torch = self.gnms_pytorch[version][variant_str] + gnm_torch, gnm_np = self._get_gnms(version, variant) parameters = { k: v[0] for k, v in self._get_default_kwargs(gnm_np, device=self.device).items() @@ -333,11 +312,7 @@ def test_no_batch(self, version: str, variant: Any): ) def test_omit_all_parameters(self, version: str, variant: Any): """If we omit all parameters, Torch GNM should return the template.""" - variant_str = variant.value if hasattr(variant, 'value') else variant - if variant_str not in self.gnms_np[version]: - self.skipTest(f'variant {variant_str} not supported in {version}.') - gnm_np = self.gnms_np[version][variant_str] - gnm_torch = self.gnms_pytorch[version][variant_str] + gnm_torch, gnm_np = self._get_gnms(version, variant) expected = gnm_np.template_vertex_positions actual = gnm_torch().detach().cpu().numpy() @@ -357,11 +332,7 @@ def test_variable_batch_and_omitted_parameters( parameter_count: int, ): """Exercise GNM with various batch dimensions and omitted parameters.""" - variant_str = variant.value if hasattr(variant, 'value') else variant - if variant_str not in self.gnms_np[version]: - self.skipTest(f'variant {variant_str} not supported in {version}.') - gnm_np = self.gnms_np[version][variant_str] - gnm_torch = self.gnms_pytorch[version][variant_str] + gnm_torch, gnm_np = self._get_gnms(version, variant) parameters = self._get_random_kwargs(gnm_np, batch_dims, device=self.device) # Omit some parameters. @@ -377,11 +348,7 @@ def test_variable_batch_and_omitted_parameters( def test_vertex_positions_world( self, version: str, variant: Any, n_batch: int ): - variant_str = variant.value if hasattr(variant, 'value') else variant - if variant_str not in self.gnms_np[version]: - self.skipTest(f'variant {variant_str} not supported in {version}.') - gnm_np = self.gnms_np[version][variant_str] - gnm_torch = self.gnms_pytorch[version][variant_str] + gnm_torch, gnm_np = self._get_gnms(version, variant) parameters = self._get_random_kwargs(gnm_np, [n_batch], device=self.device) @@ -418,10 +385,7 @@ def test_vertex_positions_world( ) def test_to_numpy_data_dict(self, version: str, variant: Any): """Tests to_numpy_data_dict method.""" - variant_str = variant.value if hasattr(variant, 'value') else variant - if variant_str not in self.gnms_np[version]: - self.skipTest(f'variant {variant_str} not supported in {version}.') - gnm = self.gnms_pytorch[version][variant_str] + gnm = self.gnms_pytorch.get_or_skip(version, variant) data_dict = gnm.to_numpy_data_dict() @@ -444,19 +408,14 @@ def test_to_numpy_data_dict(self, version: str, variant: Any): ) def test_from_gnm(self, version: str, variant: Any): """Tests from_gnm factory method.""" - variant_str = variant.value if hasattr(variant, 'value') else variant - if variant_str not in self.gnms_np[version]: - self.skipTest(f'variant {variant_str} not supported in {version}.') - gnm = self.gnms_pytorch[version][variant_str] + gnm, gnm_np = self._get_gnms(version, variant) new_gnm = gnm_pytorch.GNM.from_gnm(gnm) self.assertEqual(new_gnm.version, gnm.version) self.assertEqual(new_gnm.variant, gnm.variant) - parameters = self._get_random_kwargs( - self.gnms_np[version][variant_str], batch=[10], device=self.device - ) + parameters = self._get_random_kwargs(gnm_np, batch=[10], device=self.device) vertices_orig = gnm(**parameters) vertices_new = new_gnm(**parameters) diff --git a/gnm/shape/gnm_tensorflow_test.py b/gnm/shape/gnm_tensorflow_test.py index 17d38e82..592bfe66 100644 --- a/gnm/shape/gnm_tensorflow_test.py +++ b/gnm/shape/gnm_tensorflow_test.py @@ -14,6 +14,7 @@ """Tests for GNM TensorFlow implementation.""" +import copy import functools import itertools from typing import Any @@ -72,17 +73,22 @@ def setUpClass(cls): ) cls.addClassCleanup(cls.gnms_tf.clear) + def _get_gnms( + self, version: str, variant: Any + ) -> tuple[gnm_tensorflow.GNM, gnm_numpy.GNM]: + """Returns the TF and NumPy GNMs, or skips the test if unavailable.""" + return ( + self.gnms_tf.get_or_skip(version, variant), + self.gnms_np.get_or_skip(version, variant), + ) + @parameterized.product( version=_MAINTAINED_MAJOR_GNM_VERSIONS, variant=tuple(_SUPPORTED_VARIANTS), ) def test_properties_match(self, version: str, variant: Any): """Tests that important properties match between implementations.""" - variant_str = variant.value if hasattr(variant, 'value') else variant - if variant_str not in self.gnms_np[version]: - self.skipTest(f'variant {variant_str} not supported in {version}.') - gnm_tf = self.gnms_tf[version][variant_str] - gnm_np = self.gnms_np[version][variant_str] + gnm_tf, gnm_np = self._get_gnms(version, variant) self.assertEqual(gnm_tf.version, gnm_np.version) self.assertEqual(gnm_tf.identity_dim, gnm_np.identity_dim) @@ -99,10 +105,7 @@ def test_parity_with_gnm_numpy( self, version: str, variant: str, batch_dims: tuple[int, ...] ): """Tests that TensorFlow GNM poses vertices the same as NumPy GNM.""" - if variant not in self.gnms_np[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms_np[version][variant] - gnm_tf = self.gnms_tf[version][variant] + gnm_tf, gnm_np = self._get_gnms(version, variant) # Build a batch of random parameters. parameters_np = gnm_test_utils.random_gnm_parameters( @@ -126,10 +129,7 @@ def test_compute_vertex_normals_numpy_parity( self, version: str, variant: str, batch_dims: tuple[int, ...] ): """Tests that TensorFlow vertex normals match the NumPy implementation.""" - if variant not in self.gnms_np[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms_np[version][variant] - gnm_tf = self.gnms_tf[version][variant] + gnm_tf, gnm_np = self._get_gnms(version, variant) parameters_np = gnm_test_utils.random_gnm_parameters( gnm_np, batch_shape=batch_dims, seed=self.rng @@ -152,11 +152,7 @@ def test_vertices_and_landmarks( self, version: str, variant: Any, batch_size: tuple[int, ...] ): """Tests extracting vertices and landmarks in TensorFlow.""" - variant_str = variant.value if hasattr(variant, 'value') else variant - if variant_str not in self.gnms_np[version]: - self.skipTest(f'variant {variant_str} not supported in {version}.') - gnm_tf = self.gnms_tf[version][variant_str] - gnm_np = self.gnms_np[version][variant_str] + gnm_tf, gnm_np = self._get_gnms(version, variant) parameters_np = gnm_test_utils.random_gnm_parameters( gnm_np, batch_shape=batch_size, seed=self.rng @@ -179,10 +175,7 @@ def test_vertices_and_landmarks( def test_tf_function_behavior( self, version: str, variant: str, batch_dims: tuple[int, ...] ): - if variant not in self.gnms_np[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms_np[version][variant] - gnm_tf = self.gnms_tf[version][variant] + gnm_tf, gnm_np = self._get_gnms(version, variant) # Build a batch of random parameters. gnm_tf_parameters = gnm_test_utils.random_gnm_parameters( @@ -215,10 +208,7 @@ def tf_function_wrapper_xla(identity, expression, rotations, translation): ) def test_bad_shape(self, version: str, variant: str): """Badly shaped parameter should throw an error.""" - if variant not in self.gnms_np[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms_np[version][variant] - gnm_tf = self.gnms_tf[version][variant] + gnm_tf, gnm_np = self._get_gnms(version, variant) bad_dimension = ( gnm_tf.expression_dim + gnm_tf.identity_dim + gnm_np.num_joints @@ -238,10 +228,7 @@ def test_bad_shape(self, version: str, variant: str): ) def test_bad_shape_joint_transforms(self, version: str, variant: str): """Badly shaped parameter should throw an error.""" - if variant not in self.gnms_np[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms_np[version][variant] - gnm_tf = self.gnms_tf[version][variant] + gnm_tf, gnm_np = self._get_gnms(version, variant) bad_dimension = ( gnm_tf.expression_dim + gnm_tf.identity_dim + gnm_np.num_joints @@ -274,10 +261,7 @@ def test_joint_transforms_numpy_parity( self, version: str, variant: str, batch_dims: tuple[int, ...] ): """Test that the joint transformations function matches NumPy.""" - if variant not in self.gnms_np[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms_np[version][variant] - gnm_tf = self.gnms_tf[version][variant] + gnm_tf, gnm_np = self._get_gnms(version, variant) # Build a batch of random parameters. parameters_np = gnm_test_utils.random_gnm_parameters( @@ -317,12 +301,11 @@ def test_joint_transforms_numpy_parity( variant=tuple(_SUPPORTED_VARIANTS), ) def test_prune_vertices(self, version: str, variant: str): - if variant not in self.gnms_np[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms_np[version][variant] - gnm_tf = self.gnms_tf[version][variant] + gnm_tf, gnm_np = self._get_gnms(version, variant) - gnm_pruned = gnm_test_utils.load_gnm(gnm_tensorflow.GNM, version, variant) + # Prune a deep copy to avoid modifying the cached model for other tests: a + # shallow copy would share its landmark and topology caches. + gnm_pruned = copy.deepcopy(gnm_tf) keep_vertices = gnm_np.quads[0] gnm_pruned.prune_vertices(keep_vertices) @@ -360,10 +343,7 @@ def test_prune_vertices(self, version: str, variant: str): ) def test_no_batch(self, version: str, variant: str): """Tests we can use TF GNM without a leading batch dimension.""" - if variant not in self.gnms_np[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms_np[version][variant] - gnm_tf = self.gnms_tf[version][variant] + gnm_tf, gnm_np = self._get_gnms(version, variant) parameters = { k: v[0] for k, v in gnm_test_utils.random_gnm_parameters( @@ -379,10 +359,7 @@ def test_no_batch(self, version: str, variant: str): ) def test_omit_all_parameters(self, version: str, variant: str): """If we omit all parameters, TF GNM should return the template.""" - if variant not in self.gnms_np[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms_np[version][variant] - gnm_tf = self.gnms_tf[version][variant] + gnm_tf, gnm_np = self._get_gnms(version, variant) expected = gnm_np.template_vertex_positions actual = gnm_tf().numpy() @@ -402,10 +379,7 @@ def test_variable_batch_and_omitted_parameters( parameter_count: int, ): """Exercise GNM with various batch dimensions and omitted parameters.""" - if variant not in self.gnms_np[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms_np[version][variant] - gnm_tf = self.gnms_tf[version][variant] + gnm_tf, gnm_np = self._get_gnms(version, variant) parameters = gnm_test_utils.random_gnm_parameters( gnm_np, batch_shape=batch_dims, seed=self.rng @@ -423,10 +397,7 @@ def test_variable_batch_and_omitted_parameters( def test_vertex_positions_world( self, version: str, variant: str, batch_dims: tuple[int, ...] ): - if variant not in self.gnms_np[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm_np = self.gnms_np[version][variant] - gnm_tf = self.gnms_tf[version][variant] + gnm_tf, gnm_np = self._get_gnms(version, variant) parameters = gnm_test_utils.random_gnm_parameters( gnm_np, batch_shape=batch_dims, seed=self.rng @@ -463,10 +434,7 @@ def test_vertex_positions_world( ) def test_to_numpy_data_dict(self, version: str, variant: Any): """Tests to_numpy_data_dict method.""" - variant_str = variant.value if hasattr(variant, 'value') else variant - if variant_str not in self.gnms_np[version]: - self.skipTest(f'variant {variant_str} not supported in {version}.') - gnm = self.gnms_tf[version][variant_str] + gnm = self.gnms_tf.get_or_skip(version, variant) data_dict = gnm.to_numpy_data_dict() @@ -489,19 +457,14 @@ def test_to_numpy_data_dict(self, version: str, variant: Any): ) def test_from_gnm(self, version: str, variant: Any): """Tests from_gnm factory method.""" - variant_str = variant.value if hasattr(variant, 'value') else variant - if variant_str not in self.gnms_np[version]: - self.skipTest(f'variant {variant_str} not supported in {version}.') - gnm = self.gnms_tf[version][variant_str] + gnm, gnm_np = self._get_gnms(version, variant) new_gnm = gnm_tensorflow.GNM.from_gnm(gnm) self.assertEqual(new_gnm.version, gnm.version) self.assertEqual(new_gnm.variant, gnm.variant) - parameters = gnm_test_utils.random_gnm_parameters( - self.gnms_np[version][variant_str], seed=self.rng - ) + parameters = gnm_test_utils.random_gnm_parameters(gnm_np, seed=self.rng) parameters_tf = { k: tf.convert_to_tensor(v, dtype=tf.float32) for k, v in parameters.items() diff --git a/gnm/shape/gnm_utils_test.py b/gnm/shape/gnm_utils_test.py index ff8237e6..676f6024 100644 --- a/gnm/shape/gnm_utils_test.py +++ b/gnm/shape/gnm_utils_test.py @@ -74,27 +74,14 @@ class GNMUtilsConversionTest(BaseGNMUtilsTest): # version rather than reloading them for each pair. max_loaded = None - @parameterized.product( - version=_MAINTAINED_MAJOR_GNM_VERSIONS, - from_variant=tuple(_SUPPORTED_VARIANTS), - to_variant=tuple(_SUPPORTED_VARIANTS), - basis_type=( - gnm_utils.BasisType.IDENTITY, - gnm_utils.BasisType.EXPRESSION, - ), - ) - def test_convert_coefficients( - self, version, from_variant, to_variant, basis_type - ): - if from_variant not in self.gnms[version]: - self.skipTest( - f'from_variant {from_variant} not supported in GNM version {version}.' - ) - if to_variant not in self.gnms[version]: - self.skipTest( - f'to_variant {to_variant} not supported in GNM version {version}.' - ) - + def _get_gnms( + self, + version: str, + from_variant: gnm_numpy.GNMVariant, + to_variant: gnm_numpy.GNMVariant, + basis_type: gnm_utils.BasisType, + ) -> tuple[gnm_numpy.GNM, gnm_numpy.GNM]: + """Returns the GNMs to convert between, or skips unsupported conversions.""" # For expression, all combinations are allowed, but not for identity. if basis_type == gnm_utils.BasisType.IDENTITY: from_head = from_variant in _HEAD_VARIANTS @@ -109,8 +96,26 @@ def test_convert_coefficients( if (from_upper and to_body) or (to_upper and from_body): self.skipTest('Cannot convert identity between upper body and body.') - from_gnm = self.gnms[version][from_variant] - to_gnm = self.gnms[version][to_variant] + return ( + self.gnms.get_or_skip(version, from_variant), + self.gnms.get_or_skip(version, to_variant), + ) + + @parameterized.product( + version=_MAINTAINED_MAJOR_GNM_VERSIONS, + from_variant=tuple(_SUPPORTED_VARIANTS), + to_variant=tuple(_SUPPORTED_VARIANTS), + basis_type=( + gnm_utils.BasisType.IDENTITY, + gnm_utils.BasisType.EXPRESSION, + ), + ) + def test_convert_coefficients( + self, version, from_variant, to_variant, basis_type + ): + from_gnm, to_gnm = self._get_gnms( + version, from_variant, to_variant, basis_type + ) from_dim = gnm_utils.get_gnm_attribute( from_gnm, gnm_utils.BASIS_DIM_ATTRIBUTE_MAP[basis_type] @@ -168,27 +173,9 @@ def test_convert_coefficients( def test_raises_on_dimension_mismatch( self, version, from_variant, to_variant, basis_type ): - if from_variant not in self.gnms[version]: - self.skipTest(f'from_variant {from_variant} not supported in {version}.') - if to_variant not in self.gnms[version]: - self.skipTest(f'to_variant {to_variant} not supported in {version}.') - - # For expression, all combinations are allowed, but not for identity. - if basis_type == gnm_utils.BasisType.IDENTITY: - from_head = from_variant in _HEAD_VARIANTS - to_head = to_variant in _HEAD_VARIANTS - from_body = from_variant in _BODY_VARIANTS - to_body = to_variant in _BODY_VARIANTS - from_upper = from_variant in _UPPER_BODY_VARIANTS - to_upper = to_variant in _UPPER_BODY_VARIANTS - - if (from_head and to_body) or (to_head and from_body): - self.skipTest('Cannot convert identity between head and body.') - if (from_upper and to_body) or (to_upper and from_body): - self.skipTest('Cannot convert identity between upper body and body.') - - from_gnm = self.gnms[version][from_variant] - to_gnm = self.gnms[version][to_variant] + from_gnm, to_gnm = self._get_gnms( + version, from_variant, to_variant, basis_type + ) from_dim = gnm_utils.get_gnm_attribute( from_gnm, gnm_utils.BASIS_DIM_ATTRIBUTE_MAP[basis_type] @@ -225,9 +212,7 @@ def _random_rotations(self, gnm, batch: tuple[int, ...] = ()) -> np.ndarray: batch_shape=(tuple(), (5,), (5, 4), (5, 4, 3)), ) def test_identity_to_regions_and_back(self, version, variant, batch_shape): - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) orig_identity = self._random_identity(gnm, batch_shape) regions = gnm_utils.identity_to_regions(orig_identity, gnm) identity_again = gnm_utils.regions_to_identity(regions, gnm) @@ -239,9 +224,7 @@ def test_identity_to_regions_and_back(self, version, variant, batch_shape): batch_shape=(tuple(), (5,), (5, 4), (5, 4, 3)), ) def test_expression_to_regions_and_back(self, version, variant, batch_shape): - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) orig_expression = self._random_expression(gnm, batch_shape) regions = gnm_utils.expression_to_regions(orig_expression, gnm) expression_again = gnm_utils.regions_to_expression(regions, gnm) @@ -257,9 +240,7 @@ def test_expression_to_regions_and_back(self, version, variant, batch_shape): def test_normalize_denormalize_expression( self, version, variant, batch_shape ): - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) coeffs = self._random_expression(gnm, batch_shape) norm = gnm_utils.normalize_expression(coeffs, gnm) denorm = gnm_utils.denormalize_expression(norm, gnm) @@ -271,10 +252,7 @@ def test_normalize_denormalize_expression( batch_shape=(tuple(), (5,), (5, 4), (5, 4, 3)), ) def test_joint_rotations_cycle(self, version, variant, batch_shape): - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) rots = self._random_rotations(gnm, batch_shape) regions = gnm_utils.joint_rotations_to_regions(rots, gnm) @@ -290,9 +268,7 @@ def test_joint_rotations_cycle(self, version, variant, batch_shape): def test_identity_composed_from_individual_regions( self, version, variant, batch_shape ): - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) orig_identity = self._random_identity(gnm, batch_shape) regions = gnm_utils.identity_to_regions(orig_identity, gnm) @@ -313,9 +289,7 @@ def test_identity_composed_from_individual_regions( def test_expression_composed_from_individual_regions( self, version, variant, batch_shape ): - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) orig_expression = self._random_expression(gnm, batch_shape) regions = gnm_utils.expression_to_regions(orig_expression, gnm) @@ -335,9 +309,7 @@ def test_expression_composed_from_individual_regions( variant=tuple(_SUPPORTED_VARIANTS), ) def test_raises_on_wrong_coeffs_shape(self, version, variant): - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) coeffs = self._random_identity(gnm, (2, 3)) dim = coeffs.shape[-1] with self.assertRaisesRegex( @@ -350,9 +322,7 @@ def test_raises_on_wrong_coeffs_shape(self, version, variant): variant=tuple(_SUPPORTED_VARIANTS), ) def test_raises_on_wrong_region_name(self, version, variant): - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) bad_region_name = 'non-existent' bad_regions = {f'{bad_region_name}': np.zeros((10,))} with self.assertRaisesRegex(ValueError, f'No region {bad_region_name} in'): @@ -363,9 +333,7 @@ def test_raises_on_wrong_region_name(self, version, variant): variant=tuple(_SUPPORTED_VARIANTS), ) def test_region_identity_components_combine(self, version, variant): - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) regions = gnm_utils.region_identity_components(gnm) regions = {k: v.transpose(1, 2, 0) for k, v in regions.items()} @@ -379,9 +347,7 @@ def test_region_identity_components_combine(self, version, variant): variant=tuple(_SUPPORTED_VARIANTS), ) def test_region_expression_components_combine(self, version, variant): - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) regions = gnm_utils.region_expression_components(gnm) regions = {k: v.transpose(1, 2, 0) for k, v in regions.items()} @@ -397,9 +363,7 @@ def test_region_expression_components_combine(self, version, variant): def test_expression_region_indices_matches_expression_regions( self, version, variant ): - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) expression_vector = self._random_expression(gnm, ()) expression_regions = gnm_utils.expression_to_regions(expression_vector, gnm) expression_region_indices = gnm_utils.expression_regions_indices(gnm) @@ -417,9 +381,7 @@ def test_expression_region_indices_matches_expression_regions( def test_identity_region_indices_matches_identity_regions( self, version, variant ): - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) identity_vector = self._random_identity(gnm, ()) identity_regions = gnm_utils.identity_to_regions(identity_vector, gnm) identity_region_indices = gnm_utils.identity_regions_indices(gnm) @@ -435,9 +397,7 @@ def test_identity_region_indices_matches_identity_regions( variant=tuple(_SUPPORTED_VARIANTS), ) def test_identity_dim_is_correct(self, version, variant): - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) expected_dim = gnm.identity_dim self.assertEqual(gnm_utils.identity_dim(gnm), expected_dim) @@ -446,9 +406,7 @@ def test_identity_dim_is_correct(self, version, variant): variant=tuple(_SUPPORTED_VARIANTS), ) def test_expression_dim_is_correct(self, version, variant): - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) expected_dim = gnm.expression_dim self.assertEqual(gnm_utils.expression_dim(gnm), expected_dim) @@ -457,9 +415,7 @@ def test_expression_dim_is_correct(self, version, variant): variant=tuple(_SUPPORTED_VARIANTS), ) def test_expression_region_dims_are_correct(self, version, variant): - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) expected_dim = gnm.expression_dim regions = gnm_utils.expression_to_regions(np.zeros(expected_dim), gnm) @@ -475,9 +431,7 @@ def test_expression_region_dims_are_correct(self, version, variant): variant=tuple(_SUPPORTED_VARIANTS), ) def test_identity_region_dims_are_correct(self, version, variant): - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) expected_dim = gnm.identity_dim regions = gnm_utils.identity_to_regions(np.zeros(expected_dim), gnm) @@ -495,9 +449,7 @@ def test_identity_region_dims_are_correct(self, version, variant): ), ) def test_expression_sigmas(self, version, variant): - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) region_name = 'left_eye_region' vertex_group_name = 'expression_basis_left_eye' @@ -521,9 +473,7 @@ def test_expression_sigmas(self, version, variant): variant=tuple(_SUPPORTED_VARIANTS), ) def test_identity_region_names(self, version, variant): - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) identity = self._random_identity(gnm, ()) region_names = gnm_utils.identity_region_names(gnm) @@ -537,9 +487,7 @@ def test_identity_region_names(self, version, variant): variant=tuple(_SUPPORTED_VARIANTS), ) def test_expression_region_names(self, version, variant): - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) expression = self._random_expression(gnm, ()) region_names = gnm_utils.expression_region_names(gnm) @@ -557,9 +505,7 @@ def test_expression_region_names(self, version, variant): def test_joint_rotations_to_regions_with_missing_regions( self, version, variant, batch_shape, missing_joint_index ): - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) # if not hasattr(gnm, 'joint_names') or not gnm.joint_names: # self.skipTest(f'variant {variant} does not have joint rotations.') if missing_joint_index >= len(gnm.joint_names): @@ -584,9 +530,7 @@ def test_joint_rotations_to_regions_with_missing_regions( batch_shape=(tuple(), (5,), (5, 4), (5, 4, 3)), ) def test_scaling_for_expression_loss(self, version, variant, batch_shape): - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) expression1 = self._random_expression(gnm, batch_shape) expression2 = self._random_expression(gnm, batch_shape) @@ -611,9 +555,7 @@ def test_scaling_for_expression_loss(self, version, variant, batch_shape): batch_shape=(tuple(), (5,), (5, 4), (5, 4, 3)), ) def test_scaling_for_identity_loss(self, version, variant, batch_shape): - if variant not in self.gnms[version]: - self.skipTest(f'variant {variant} not supported in {version}.') - gnm = self.gnms[version][variant] + gnm = self.gnms.get_or_skip(version, variant) identity1 = self._random_identity(gnm, batch_shape) identity2 = self._random_identity(gnm, batch_shape) diff --git a/gnm/shape/visualization/gnm_mitsuba_test.py b/gnm/shape/visualization/gnm_mitsuba_test.py index 7a1e2e4c..84c211e6 100644 --- a/gnm/shape/visualization/gnm_mitsuba_test.py +++ b/gnm/shape/visualization/gnm_mitsuba_test.py @@ -19,6 +19,7 @@ from absl.testing import absltest from absl.testing import parameterized from gnm.shape import gnm_numpy +from gnm.shape import gnm_test_utils from gnm.shape.data.versions import gnm_test_catalog from gnm.shape.visualization import gnm_mitsuba import mitsuba as mi @@ -40,9 +41,8 @@ def setUpClass(cls): if not mi.variant(): mi.set_variant('cuda_ad_rgb', 'llvm_ad_rgb') cls.gnms = { - version: gnm_numpy.GNM.from_remote( - gnm_numpy.GNMMajorVersion(version.removeprefix('v')), - gnm_numpy.GNMVariant.HEAD, + version: gnm_test_utils.load_gnm( + gnm_numpy.GNM, version, gnm_numpy.GNMVariant.HEAD ) for version in _MAINTAINED_MAJOR_GNM_VERSIONS } diff --git a/gnm/shape/visualization/gnm_pyrender_test.py b/gnm/shape/visualization/gnm_pyrender_test.py index 9ff5f1a7..df07c152 100644 --- a/gnm/shape/visualization/gnm_pyrender_test.py +++ b/gnm/shape/visualization/gnm_pyrender_test.py @@ -21,6 +21,7 @@ from absl.testing import parameterized from etils import epath from gnm.shape import gnm_numpy +from gnm.shape import gnm_test_utils from gnm.shape.data.versions import gnm_test_catalog from gnm.shape.visualization import gnm_pyrender import mediapy as media @@ -41,9 +42,8 @@ class GNMPyrenderTest(parameterized.TestCase): def setUpClass(cls): super().setUpClass() cls.gnms = { - version: gnm_numpy.GNM.from_remote( - gnm_numpy.GNMMajorVersion(version.removeprefix('v')), - gnm_numpy.GNMVariant.HEAD, + version: gnm_test_utils.load_gnm( + gnm_numpy.GNM, version, gnm_numpy.GNMVariant.HEAD ) for version in _MAINTAINED_MAJOR_GNM_VERSIONS } diff --git a/gnm/shape/visualization/render_common_test.py b/gnm/shape/visualization/render_common_test.py index e3250132..bf9901eb 100644 --- a/gnm/shape/visualization/render_common_test.py +++ b/gnm/shape/visualization/render_common_test.py @@ -14,10 +14,13 @@ """Tests for backend-agnostic helpers in render_common.""" +import functools from typing import Any + from absl.testing import absltest from absl.testing import parameterized from gnm.shape import gnm_numpy +from gnm.shape import gnm_test_utils from gnm.shape.data.versions import gnm_test_catalog from gnm.shape.visualization import camera_conversions from gnm.shape.visualization import render_common @@ -26,6 +29,17 @@ _TupleOfInts = tuple[int, ...] +@functools.cache +def _load_head_gnms() -> dict[str, gnm_numpy.GNM]: + """Loads the GNM head models once, shared by all test classes.""" + return { + version: gnm_test_utils.load_gnm( + gnm_numpy.GNM, version, gnm_numpy.GNMVariant.HEAD + ) + for version in gnm_test_catalog.MAINTAINED_MAJOR_VERSIONS + } + + def _get_random_parameters( batch_dims: _TupleOfInts, gnm_np: gnm_numpy.GNM, @@ -94,12 +108,7 @@ class TestGetLookAtWorldToCamera(parameterized.TestCase): @classmethod def setUpClass(cls): super().setUpClass() - cls.gnms = {} - for version in gnm_test_catalog.MAINTAINED_MAJOR_VERSIONS: - cls.gnms[version] = gnm_numpy.GNM.from_remote( - gnm_numpy.GNMMajorVersion(version.removeprefix('v')), - gnm_numpy.GNMVariant.HEAD, - ) + cls.gnms = _load_head_gnms() @parameterized.named_parameters(*[ (version, version) @@ -206,12 +215,7 @@ class TestGetFillFactorCameraToImage(parameterized.TestCase): @classmethod def setUpClass(cls): super().setUpClass() - cls.gnms = {} - for version in gnm_test_catalog.MAINTAINED_MAJOR_VERSIONS: - cls.gnms[version] = gnm_numpy.GNM.from_remote( - gnm_numpy.GNMMajorVersion(version.removeprefix('v')), - gnm_numpy.GNMVariant.HEAD, - ) + cls.gnms = _load_head_gnms() @parameterized.named_parameters(*[ (version, version) @@ -248,12 +252,7 @@ class TestLoadTexture(parameterized.TestCase): @classmethod def setUpClass(cls): super().setUpClass() - cls.gnms = {} - for version in gnm_test_catalog.MAINTAINED_MAJOR_VERSIONS: - cls.gnms[version] = gnm_numpy.GNM.from_remote( - gnm_numpy.GNMMajorVersion(version.removeprefix('v')), - gnm_numpy.GNMVariant.HEAD, - ) + cls.gnms = _load_head_gnms() @parameterized.named_parameters(*[ (version, version) @@ -287,12 +286,7 @@ class TestProjectPointsForGnm(parameterized.TestCase): @classmethod def setUpClass(cls): super().setUpClass() - cls.gnms = {} - for version in gnm_test_catalog.MAINTAINED_MAJOR_VERSIONS: - cls.gnms[version] = gnm_numpy.GNM.from_remote( - gnm_numpy.GNMMajorVersion(version.removeprefix('v')), - gnm_numpy.GNMVariant.HEAD, - ) + cls.gnms = _load_head_gnms() @parameterized.named_parameters(*[ (version, version) @@ -324,12 +318,7 @@ class TestGetSpinWorldToCamera(parameterized.TestCase): @classmethod def setUpClass(cls): super().setUpClass() - cls.gnms = {} - for version in gnm_test_catalog.MAINTAINED_MAJOR_VERSIONS: - cls.gnms[version] = gnm_numpy.GNM.from_remote( - gnm_numpy.GNMMajorVersion(version.removeprefix('v')), - gnm_numpy.GNMVariant.HEAD, - ) + cls.gnms = _load_head_gnms() @parameterized.named_parameters(*[ (version, version) @@ -363,12 +352,9 @@ class TestRenderGNMMesh(parameterized.TestCase): @classmethod def setUpClass(cls): super().setUpClass() - cls.gnm_np = gnm_numpy.GNM.from_remote( - gnm_numpy.GNMMajorVersion( - gnm_test_catalog.MAINTAINED_MAJOR_VERSIONS[0].removeprefix('v') - ), - gnm_numpy.GNMVariant.HEAD, - ) + cls.gnm_np = _load_head_gnms()[ + gnm_test_catalog.MAINTAINED_MAJOR_VERSIONS[0] + ] def test_render_gnm_mesh_mock(self): """Tests render_gnm_mesh prepares parameters and calls backend_render_fn.""" diff --git a/gnm/shape/visualization/render_gnm_parity_test.py b/gnm/shape/visualization/render_gnm_parity_test.py index 79be0b4f..9c78537a 100644 --- a/gnm/shape/visualization/render_gnm_parity_test.py +++ b/gnm/shape/visualization/render_gnm_parity_test.py @@ -18,7 +18,6 @@ from absl.testing import absltest from absl.testing import parameterized from etils import epath -from gnm.shape import gnm_numpy from gnm.shape.data.versions import gnm_test_catalog from gnm.shape.visualization import render_gnm from gnm.shape.visualization import render_gnm_mitsuba @@ -34,11 +33,8 @@ class RenderGNMParityTest(parameterized.TestCase): def setUp(self): super().setUp() - self.gnm_np = gnm_numpy.GNM.from_remote( - gnm_numpy.GNMMajorVersion( - gnm_test_catalog.MAINTAINED_MAJOR_VERSIONS[0].removeprefix('v') - ), - gnm_numpy.GNMVariant.HEAD, + self.gnm_np = render_gnm_test_base.load_head_gnm( + gnm_test_catalog.MAINTAINED_MAJOR_VERSIONS[0] ) self.outputs_dir = _OUTPUTS_DIR / self.__class__.__name__ self.outputs_dir.mkdir(parents=True, exist_ok=True) diff --git a/gnm/shape/visualization/render_gnm_test.py b/gnm/shape/visualization/render_gnm_test.py index 6d83450d..7ada8ab1 100644 --- a/gnm/shape/visualization/render_gnm_test.py +++ b/gnm/shape/visualization/render_gnm_test.py @@ -53,12 +53,10 @@ class TestProjectPointsForGNM(parameterized.TestCase): @classmethod def setUpClass(cls): super().setUpClass() - cls.gnms = {} - for version in gnm_test_catalog.MAINTAINED_MAJOR_VERSIONS: - cls.gnms[version] = gnm_numpy.GNM.from_remote( - gnm_numpy.GNMMajorVersion(version.removeprefix('v')), - gnm_numpy.GNMVariant.HEAD, - ) + cls.gnms = { + version: render_gnm_test_base.load_head_gnm(version) + for version in gnm_test_catalog.MAINTAINED_MAJOR_VERSIONS + } def setUp(self): super().setUp() diff --git a/gnm/shape/visualization/render_gnm_test_base.py b/gnm/shape/visualization/render_gnm_test_base.py index 4f092ac5..dfafc8ad 100644 --- a/gnm/shape/visualization/render_gnm_test_base.py +++ b/gnm/shape/visualization/render_gnm_test_base.py @@ -15,6 +15,7 @@ """Shared base test classes and helpers for GNM mesh rendering backends.""" from collections.abc import Callable +import functools import tempfile from typing import Any @@ -24,6 +25,7 @@ import cv2 from etils import epath from gnm.shape import gnm_numpy +from gnm.shape import gnm_test_utils from gnm.shape.data.versions import gnm_test_catalog from gnm.shape.visualization import render_common from gnm.shape.visualization import vertex_colors as vertex_colors_module @@ -34,6 +36,23 @@ _OUTPUTS_DIR = epath.Path(_OUTPUTS_TMPDIR.name) +@functools.cache +def load_head_gnm(version: str) -> gnm_numpy.GNM: + """Loads the GNM head model of `version` once, shared by all tests. + + Tests must not modify the returned model. + + Args: + version: The GNM major version, e.g. 'v3'. + + Returns: + The GNM head model. + """ + return gnm_test_utils.load_gnm( + gnm_numpy.GNM, version, gnm_numpy.GNMVariant.HEAD + ) + + def write_gif( outputs_dir: epath.Path | str, name: str, @@ -102,12 +121,10 @@ class RenderGNMTestBase(parameterized.TestCase): @classmethod def setUpClass(cls): super().setUpClass() - cls.gnms = {} - for version in gnm_test_catalog.MAINTAINED_MAJOR_VERSIONS: - cls.gnms[version] = gnm_numpy.GNM.from_remote( - gnm_numpy.GNMMajorVersion(version.removeprefix('v')), - gnm_numpy.GNMVariant.HEAD, - ) + cls.gnms = { + version: load_head_gnm(version) + for version in gnm_test_catalog.MAINTAINED_MAJOR_VERSIONS + } def setUp(self): super().setUp() @@ -549,12 +566,7 @@ def mock_render(vertices, image_size, **kwargs): ) self.mock_render.side_effect = mock_render - self.gnm_np = gnm_numpy.GNM.from_remote( - gnm_numpy.GNMMajorVersion( - gnm_test_catalog.MAINTAINED_MAJOR_VERSIONS[0].removeprefix('v') - ), - gnm_numpy.GNMVariant.HEAD, - ) + self.gnm_np = load_head_gnm(gnm_test_catalog.MAINTAINED_MAJOR_VERSIONS[0]) self.image_size = (240, 320) self.image_dims = (self.image_size[1], self.image_size[0], 3) self.vertices_shape = self.gnm_np.template_vertex_positions.shape diff --git a/gnm/shape/visualization/vertex_colors_test.py b/gnm/shape/visualization/vertex_colors_test.py index bed7d026..00aff381 100644 --- a/gnm/shape/visualization/vertex_colors_test.py +++ b/gnm/shape/visualization/vertex_colors_test.py @@ -14,9 +14,12 @@ """Tests for vertex colors visualization.""" +import functools + from absl.testing import absltest from absl.testing import parameterized from gnm.shape import gnm_numpy +from gnm.shape import gnm_test_utils from gnm.shape.data.versions import gnm_test_catalog from gnm.shape.visualization import vertex_colors import numpy as np @@ -35,15 +38,18 @@ class VertexColorsTest(parameterized.TestCase): def setUpClass(cls): super().setUpClass() - cls.models = {} - for version in _MAINTAINED_MAJOR_GNM_VERSIONS: - cls.models[version] = {} - for variant in _MAJOR_VERSION_TO_VARIANTS_MAP[version]: - if variant in _SUPPORTED_VARIANTS: - cls.models[version][variant] = gnm_numpy.GNM.from_remote( - gnm_numpy.GNMMajorVersion(version.removeprefix('v')), - gnm_numpy.GNMVariant(variant), - ) + cls.models = gnm_test_utils.LazyGNMDict( + load_fn=functools.partial(gnm_test_utils.load_gnm, gnm_numpy.GNM), + variants_by_version={ + version: [ + variant + for variant in _MAJOR_VERSION_TO_VARIANTS_MAP[version] + if variant in _SUPPORTED_VARIANTS + ] + for version in _MAINTAINED_MAJOR_GNM_VERSIONS + }, + ) + cls.addClassCleanup(cls.models.clear) @parameterized.product( version=_MAINTAINED_MAJOR_GNM_VERSIONS, @@ -51,10 +57,7 @@ def setUpClass(cls): ) def test_get_vertex_colors(self, version, variant): """Tests we can get per-vertex colors.""" - if variant not in self.models[version]: - variant_name = variant.value if hasattr(variant, 'value') else variant - self.skipTest(f'variant {variant_name} not supported in {version}.') - gnm_np = self.models[version][variant] + gnm_np = self.models.get_or_skip(version, variant) colors = vertex_colors.get_vertex_colors(gnm_np) # Sanity check on color values. @@ -65,10 +68,7 @@ def test_get_vertex_colors(self, version, variant): variant=[gnm_numpy.GNMVariant.HEAD], ) def test_gets_inner_head_colors(self, version, variant): - if variant not in self.models[version]: - variant_name = variant.value if hasattr(variant, 'value') else variant - self.skipTest(f'variant {variant_name} not supported in {version}.') - gnm_np = self.models[version][variant] + gnm_np = self.models.get_or_skip(version, variant) colors = vertex_colors.get_vertex_colors_for_inner_head(gnm_np) groups = ( ('mouth_sock',), @@ -94,4 +94,4 @@ def test_gets_inner_head_colors(self, version, variant): if __name__ == '__main__': - absltest.main() + absltest.main(testLoader=gnm_test_utils.ModelOrderedTestLoader())