From 22f1069cc8f0fe93da27fe00e7b8684bffd00f40 Mon Sep 17 00:00:00 2001
From: The TensorFlow Datasets Authors <no-reply@google.com>
Date: Mon, 21 Jul 2025 14:44:04 -0700
Subject: [PATCH] Support huggingface Video feature

PiperOrigin-RevId: 785582494
---
 .../core/features/video_feature.py            | 19 +++++++++++--------
 .../core/features/video_feature_test.py       | 17 ++++++++++++++++-
 .../core/utils/huggingface_utils.py           |  2 ++
 3 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/tensorflow_datasets/core/features/video_feature.py b/tensorflow_datasets/core/features/video_feature.py
index 75114c7d494..c9c1d5e24bc 100644
--- a/tensorflow_datasets/core/features/video_feature.py
+++ b/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):
diff --git a/tensorflow_datasets/core/features/video_feature_test.py b/tensorflow_datasets/core/features/video_feature_test.py
index 835b9452f59..7fff18e39c9 100644
--- a/tensorflow_datasets/core/features/video_feature_test.py
+++ b/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()
diff --git a/tensorflow_datasets/core/utils/huggingface_utils.py b/tensorflow_datasets/core/utils/huggingface_utils.py
index 26876ee516f..84bb3bd39ad 100644
--- a/tensorflow_datasets/core/utils/huggingface_utils.py
+++ b/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.')