diff --git a/gnm/shape/data/semantic_sampler/expression_decoder_model.h5 b/gnm/shape/data/semantic_sampler/v3_0/expression_decoder_model.h5 similarity index 100% rename from gnm/shape/data/semantic_sampler/expression_decoder_model.h5 rename to gnm/shape/data/semantic_sampler/v3_0/expression_decoder_model.h5 diff --git a/gnm/shape/data/semantic_sampler/identity_decoder_model.h5 b/gnm/shape/data/semantic_sampler/v3_0/identity_decoder_model.h5 similarity index 100% rename from gnm/shape/data/semantic_sampler/identity_decoder_model.h5 rename to gnm/shape/data/semantic_sampler/v3_0/identity_decoder_model.h5 diff --git a/gnm/shape/semantic_sampler.py b/gnm/shape/semantic_sampler.py index 35f30ffb..32e08279 100644 --- a/gnm/shape/semantic_sampler.py +++ b/gnm/shape/semantic_sampler.py @@ -38,19 +38,43 @@ from collections.abc import Mapping import enum from etils import epath +from gnm.shape.data.versions import gnm_specs +import immutabledict import numpy as np import tensorflow as tf +GNMMajorVersion = gnm_specs.GNMMajorVersion + _pkg = __package__ or 'gnm.shape' _DATA_DIR = epath.resource_path(_pkg) / 'data' -_EXPRESSION_DECODER_PATH = ( - _DATA_DIR / 'semantic_sampler' / 'expression_decoder_model.h5' -) -_IDENTITY_DECODER_PATH = ( - _DATA_DIR / 'semantic_sampler' / 'identity_decoder_model.h5' +_SEMANTIC_SAMPLER_DIR = _DATA_DIR / 'semantic_sampler' +_EXPRESSION_DECODER_FILENAME = 'expression_decoder_model.h5' +_IDENTITY_DECODER_FILENAME = 'identity_decoder_model.h5' +_VERSION_TO_DIRECTORY_MAP: Mapping[GNMMajorVersion, str] = ( + immutabledict.immutabledict({ + GNMMajorVersion.V3: 'v3_0', + }) ) +def _get_decoder_model_path( + version: GNMMajorVersion, + filename: str, +) -> epath.Path: + """Returns the default decoder model path for the given GNM major version.""" + version_directory = _VERSION_TO_DIRECTORY_MAP.get(version) + if version_directory is None: + supported_versions = ', '.join( + supported_version.name + for supported_version in _VERSION_TO_DIRECTORY_MAP + ) + raise ValueError( + f'Unsupported semantic sampler version: {version!r}. Supported' + f' versions are: {supported_versions}.' + ) + return _SEMANTIC_SAMPLER_DIR / version_directory / filename + + class Gender(enum.IntEnum): """Gender classes for identity sampling.""" @@ -102,21 +126,28 @@ class ExpressionSampler: def __init__( self, - decoder_model_path: str | None = None, + decoder_model_path: epath.PathLike | None = None, + version: GNMMajorVersion = GNMMajorVersion.V3, verbose: bool = False, ) -> None: - """Initializes the ExpressionSampler class. + """Initializes the instance. Loads the decoder model and stores parameters. Args: decoder_model_path: Path to the saved decoder model file. If None, the - default model is loaded. + default model for the specified `version` is loaded. + version: The GNM major version of the default decoder model to load when + `decoder_model_path` is None. verbose: If True, prints a summary of the loaded decoder model. + + Raises: + ValueError: If `decoder_model_path` is None and `version` is not + supported. """ if decoder_model_path is None: - decoder_model_path = ( - _EXPRESSION_DECODER_PATH # pyrefly: ignore[bad-assignment] + decoder_model_path = _get_decoder_model_path( + version, _EXPRESSION_DECODER_FILENAME ) self._decoder = tf.keras.models.load_model(str(decoder_model_path)) self._expression_names = tuple(member.name.lower() for member in Expression) @@ -338,23 +369,30 @@ class IdentitySampler: def __init__( self, - decoder_model_path: str | None = None, + decoder_model_path: epath.PathLike | None = None, + version: GNMMajorVersion = GNMMajorVersion.V3, verbose: bool = False, ) -> None: - """Initializes the IdentitySampler class. + """Initializes the instance. Loads the decoder model and stores parameters. Args: decoder_model_path: Path to the saved identity decoder model file. If - None, the default model is loaded. + None, the default model for the specified `version` is loaded. + version: The GNM major version of the default decoder model to load when + `decoder_model_path` is None. verbose: If True, prints a summary of the loaded decoder model. + + Raises: + ValueError: If `decoder_model_path` is None and `version` is not + supported. """ if decoder_model_path is None: - decoder_model_path = ( - _IDENTITY_DECODER_PATH # pyrefly: ignore[bad-assignment] + decoder_model_path = _get_decoder_model_path( + version, _IDENTITY_DECODER_FILENAME ) - self._decoder = tf.keras.models.load_model(decoder_model_path) + self._decoder = tf.keras.models.load_model(str(decoder_model_path)) self._condition_dim = self._decoder.inputs[1].shape[-1] if verbose: diff --git a/gnm/shape/semantic_sampler_test.py b/gnm/shape/semantic_sampler_test.py index 9dd62800..bb586f4a 100644 --- a/gnm/shape/semantic_sampler_test.py +++ b/gnm/shape/semantic_sampler_test.py @@ -140,6 +140,68 @@ def test_non_deterministic_sampling(self): ident2 = self.identity_sampler.sample_identity(gender, ethnicity, rng=None) self.assertFalse(np.array_equal(ident1, ident2)) + def test_sampler_initialization_with_explicit_version(self): + expression_sampler = semantic_sampler.ExpressionSampler( + version=semantic_sampler.GNMMajorVersion.V3 + ) + identity_sampler = semantic_sampler.IdentitySampler( + version=semantic_sampler.GNMMajorVersion.V3 + ) + + seed = 123 + expected_expression = self.expression_sampler.sample_expression( + semantic_sampler.Expression.HAPPY, + rng=np.random.default_rng(seed), + ) + actual_expression = expression_sampler.sample_expression( + semantic_sampler.Expression.HAPPY, + rng=np.random.default_rng(seed), + ) + np.testing.assert_array_equal(actual_expression, expected_expression) + + expected_identity = self.identity_sampler.sample_identity( + semantic_sampler.Gender.FEMALE, + semantic_sampler.Ethnicity.WHITE, + rng=np.random.default_rng(seed), + ) + actual_identity = identity_sampler.sample_identity( + semantic_sampler.Gender.FEMALE, + semantic_sampler.Ethnicity.WHITE, + rng=np.random.default_rng(seed), + ) + np.testing.assert_array_equal(actual_identity, expected_identity) + + def test_sampler_initialization_with_explicit_decoder_model_path(self): + expression_model_path = ( + semantic_sampler._SEMANTIC_SAMPLER_DIR + / 'v3_0' + / 'expression_decoder_model.h5' + ) + identity_model_path = ( + semantic_sampler._SEMANTIC_SAMPLER_DIR + / 'v3_0' + / 'identity_decoder_model.h5' + ) + + expression_sampler = semantic_sampler.ExpressionSampler( + decoder_model_path=expression_model_path + ) + identity_sampler = semantic_sampler.IdentitySampler( + decoder_model_path=identity_model_path + ) + + generated_expressions = expression_sampler.sample_expression( + semantic_sampler.Expression.HAPPY, + num_samples=1, + ) + generated_identities = identity_sampler.sample_identity( + semantic_sampler.Gender.FEMALE, + semantic_sampler.Ethnicity.WHITE, + num_samples=1, + ) + self.assertEqual(generated_expressions.shape, (1, self.expression_dim)) + self.assertEqual(generated_identities.shape, (1, self.identity_dim)) + if __name__ == '__main__': absltest.main()