Starcraft dataset: create an S3 version (non default).

PiperOrigin-RevId: 281695849

Author: pierrot0

tensorflow_datasets/video/starcraft.py

@@ -56,6 +56,8 @@ def __init__(self, map_name, resolution, size_in_gb, **kwargs):
     super(StarcraftVideoConfig, self).__init__(
         version=tfds.core.Version(
             "0.1.2", experiments={tfds.core.Experiment.S3: False}),
+        supported_versions=[tfds.core.Version(
+            "1.0.0", "New split API (https://tensorflow.org/datasets/splits)")],
         **kwargs)
     self.map_name = map_name
     self.resolution = resolution
@@ -205,17 +207,19 @@ def _parse_single_video(self, example_proto):
   def _generate_examples(self, files):
     logging.info("Reading data from %s.", ",".join(files))
     with tf.Graph().as_default():
-      ds = tf.data.TFRecordDataset(files)
+      ds = tf.data.TFRecordDataset(sorted(files))
       ds = ds.map(
           self._parse_single_video,
           num_parallel_calls=tf.data.experimental.AUTOTUNE)
       iterator = tf.compat.v1.data.make_one_shot_iterator(ds).get_next()
       with tf.compat.v1.Session() as sess:
         sess.run(tf.compat.v1.global_variables_initializer())
         try:
+          i = 0
           while True:
             video = sess.run(iterator)
-            yield {"rgb_screen": video}
+            yield i, {"rgb_screen": video}
+            i += 1
 
         except tf.errors.OutOfRangeError:
           # End of file.

tensorflow_datasets/video/starcraft_test.py

@@ -60,5 +60,9 @@ class StarcraftVideoDataset128Test(testing.DatasetBuilderTestCase):
   }
 
 
+class StarcraftVideoDataset128S3Test(StarcraftVideoDataset128Test):
+  VERSION = "experimental_latest"
+
+
 if __name__ == "__main__":
   testing.test_main()