Support huggingface Video feature

PiperOrigin-RevId: 785582494

Author: The TensorFlow Datasets Authors

tensorflow_datasets/core/features/video_feature.py

@@ -93,7 +93,7 @@ class Video(sequence_feature.Sequence):
 
   def __init__(
       self,
-      shape: Sequence[Optional[int]],
+      shape: Sequence[Optional[int]] | None = None,
       encoding_format: str = 'png',
       ffmpeg_extra_args: Sequence[str] = (),
       use_colormap: bool = False,
@@ -103,8 +103,8 @@ def __init__(
     """Initializes the connector.
 
     Args:
-      shape: tuple of ints, the shape of the video (num_frames, height, width,
-        channels), where channels is 1 or 3.
+      shape: The shape of the video (num_frames, height, width, channels), where
+        channels is 1 or 3.
       encoding_format: The video is stored as a sequence of encoded images. You
         can use any encoding format supported by image_feature.Feature.
       ffmpeg_extra_args: A sequence of additional args to be passed to the
@@ -121,19 +121,22 @@ def __init__(
       ValueError: If the shape is invalid
     """
     dtype = tf.dtypes.as_dtype(dtype)
-    shape = tuple(shape)
-    if len(shape) != 4:
-      raise ValueError('Video shape should be of rank 4')
+    frame_shape = None
+    if shape:
+      shape = tuple(shape)
+      if len(shape) != 4:
+        raise ValueError('Video shape should be of rank 4')
+      frame_shape = shape[1:]
     self._encoding_format = encoding_format
     self._extra_ffmpeg_args = list(ffmpeg_extra_args or [])
     super(Video, self).__init__(
         image_feature.Image(
-            shape=shape[1:],
+            shape=frame_shape,
             dtype=dtype,
             encoding_format=encoding_format,
             use_colormap=use_colormap,
         ),
-        length=shape[0],
+        length=shape[0] if shape else None,
     )
 
   def _ffmpeg_decode(self, path_or_fobj):

tensorflow_datasets/core/features/video_feature_test.py

@@ -48,6 +48,22 @@ def test_video_numpy(self):
         test_attributes=dict(_encoding_format='png', _extra_ffmpeg_args=[]),
     )
 
+  def test_video_with_none_shape(self):
+    np_video = np.random.randint(256, size=(128, 64, 64, 3), dtype=np.uint8)
+
+    self.assertFeature(
+        feature=features.Video(shape=None),
+        shape=(None, None, None, 3),
+        dtype=tf.uint8,
+        tests=[
+            testing.FeatureExpectationItem(
+                value=np_video,
+                expected=np_video,
+            ),
+        ],
+        test_attributes=dict(_encoding_format='png', _extra_ffmpeg_args=[]),
+    )
+
   def test_video_concatenated_frames(self):
     video_shape = (None, 400, 640, 3)
     lsun_examples_path = os.path.join(self._test_data_path, 'lsun_examples')
@@ -119,6 +135,5 @@ def read(self, *args, **kwargs):
           ],
       )
 
-
 if __name__ == '__main__':
   testing.test_main()

tensorflow_datasets/core/utils/huggingface_utils.py

@@ -119,6 +119,8 @@ def convert_hf_features(hf_features) -> feature_lib.FeatureConnector:
           sample_rate=hf_features.sampling_rate,
           dtype=np.int32,
       )
+    case hf_datasets.Video():
+      return feature_lib.Video()
 
   raise TypeError(f'Type {type(hf_features)} is not supported.')