diff --git a/tensorflow_datasets/core/features/audio_feature.py b/tensorflow_datasets/core/features/audio_feature.py index 432d48fc5b0..68180ccedf1 100644 --- a/tensorflow_datasets/core/features/audio_feature.py +++ b/tensorflow_datasets/core/features/audio_feature.py @@ -82,10 +82,14 @@ def __init__( """ self._file_format = file_format self._np_dtype = np_dtype - self._dtype = tf.dtypes.as_dtype(self._np_dtype) self._shape = shape self._channels = shape[1] if len(shape) > 1 else 1 + @property + def _dtype(self) -> tf.dtypes.DType: + # Computed on access so that defining the feature doesn't import TensorFlow. + return tf.dtypes.as_dtype(self._np_dtype) + @abc.abstractmethod def encode_audio( self, fobj: BinaryIO, file_format: Optional[str] diff --git a/tensorflow_datasets/core/features/audio_feature_test.py b/tensorflow_datasets/core/features/audio_feature_test.py index ecb5cf8ed6e..a90c51f311b 100644 --- a/tensorflow_datasets/core/features/audio_feature_test.py +++ b/tensorflow_datasets/core/features/audio_feature_test.py @@ -18,6 +18,7 @@ import array import pathlib import tempfile +from unittest import mock from absl.testing import parameterized import numpy as np @@ -25,6 +26,7 @@ import tensorflow as tf from tensorflow_datasets import testing from tensorflow_datasets.core import features +from tensorflow_datasets.core.features import audio_feature class AudioFeatureTest( @@ -190,6 +192,14 @@ def read(self, *args, **kwargs): ], ) + def test_construction_does_not_use_tensorflow(self): + # Defining the feature must not import TensorFlow, so TFDS can be used + # without it (e.g. with `tfds.data_source`). + with mock.patch.object(audio_feature, 'tf') as tf_mock: + features.Audio(sample_rate=1000) + features.Audio(sample_rate=1000, file_format='wav', shape=(None, 2)) + tf_mock.dtypes.as_dtype.assert_not_called() + def test_sample_rate_property(self): self.assertEqual(features.Audio(sample_rate=1600).sample_rate, 1600)