Add MovingMNIST dataset

tensorflow_datasets/core/features/sequence_feature.py

@@ -260,7 +260,7 @@ def np_to_list(elem):
     return elem
   elif isinstance(elem, np.ndarray):
     elem = np.split(elem, elem.shape[0])
-    elem = np.squeeze(elem, axis=0)
+    elem = [np.squeeze(e, axis=0) for e in elem]
     return elem
   else:
     raise ValueError(

tensorflow_datasets/core/features/video_feature.py

@@ -29,10 +29,11 @@ class Video(sequence_feature.Sequence):
   """`FeatureConnector` for videos, png-encoding frames on disk.
 
   Video: The image connector accepts as input:
-    * uint8 array representing an video.
+    * uint8 array representing a video.
 
   Output:
-    video: tf.Tensor of type tf.uint8 and shape [num_frames, height, width, 3]
+    video: tf.Tensor of type tf.uint8 and shape
+      [num_frames, height, width, channels], where channels must be 1 or 3
 
   Example:
     * In the DatasetInfo object:
@@ -51,7 +52,7 @@ def __init__(self, shape):
 
     Args:
       shape: tuple of ints, the shape of the video (num_frames, height, width,
-        channels=3).
+        channels), where channels is 1 or 3.
 
     Raises:
       ValueError: If the shape is invalid
@@ -61,6 +62,8 @@ def __init__(self, shape):
       raise ValueError('Video shape should be of rank 4')
     if shape.count(None) > 1:
       raise ValueError('Video shape cannot have more than 1 unknown dim')
+    if shape[-1] not in (1, 3):
+      raise ValueError('Video channels must be 1 or 3, got %d' % shape[-1])
 
     super(Video, self).__init__(
         image_feature.Image(shape=shape[1:], encoding_format='png'),

tensorflow_datasets/video/__init__.py

@@ -16,5 +16,6 @@
 """Video datasets."""
 
 from tensorflow_datasets.video.bair_robot_pushing import BairRobotPushingSmall
+from tensorflow_datasets.video.moving_mnist import MovingMnist
 from tensorflow_datasets.video.starcraft import StarcraftVideo
 from tensorflow_datasets.video.starcraft import StarcraftVideoConfig

tensorflow_datasets/video/moving_mnist.py

@@ -0,0 +1,80 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import collections
+import functools
+import numpy as np
+import tensorflow as tf
+import tensorflow_datasets.public_api as tfds
+
+_OUT_RESOLUTION = (64, 64)
+_SEQUENCE_LENGTH = 20
+_citation = """
+@article{DBLP:journals/corr/SrivastavaMS15,
+  author    = {Nitish Srivastava and
+               Elman Mansimov and
+               Ruslan Salakhutdinov},
+  title     = {Unsupervised Learning of Video Representations using LSTMs},
+  journal   = {CoRR},
+  volume    = {abs/1502.04681},
+  year      = {2015},
+  url       = {http://arxiv.org/abs/1502.04681},
+  archivePrefix = {arXiv},
+  eprint    = {1502.04681},
+  timestamp = {Mon, 13 Aug 2018 16:47:05 +0200},
+  biburl    = {https://dblp.org/rec/bib/journals/corr/SrivastavaMS15},
+  bibsource = {dblp computer science bibliography, https://dblp.org}
+}
+"""
+
+
+class MovingMnist(tfds.core.GeneratorBasedBuilder):
+
+  VERSION = tfds.core.Version("0.1.0")
+
+  def _info(self):
+    return tfds.core.DatasetInfo(
+        builder=self,
+        description=(
+          "Moving variant of MNIST database of handwritten digits. This is the "
+          "data used by the authors for reporting model performance. See "
+          "`tfds.video.moving_sequence` for functions to generate training/"
+          "validation data."),
+        features=tfds.features.FeaturesDict(
+          dict(image_sequence=tfds.features.Video(
+              shape=(_SEQUENCE_LENGTH,) + _OUT_RESOLUTION + (1,)))),
+        urls=["http://www.cs.toronto.edu/~nitish/unsupervised_video/"],
+        citation=_citation,
+        splits=[tfds.Split.TEST]
+    )
+
+  def _split_generators(self, dl_manager):
+    data_path = dl_manager.download(
+        "http://www.cs.toronto.edu/~nitish/unsupervised_video/"
+        "mnist_test_seq.npy")
+
+    # authors only provide test data. See `tfds.video.moving_sequence` for
+    # approach based on creating sequences from existing datasets
+    return [
+        tfds.core.SplitGenerator(
+            name=tfds.Split.TEST,
+            num_shards=5,
+            gen_kwargs=dict(data_path=data_path)),
+    ]
+
+  def _generate_examples(self, data_path):
+    """Generate MOVING_MNIST sequences.
+
+    Args:
+      data_path (str): Path to the data file
+
+    Yields:
+      20 x 64 x 64 x 1 uint8 numpy arrays
+    """
+    with tf.io.gfile.GFile(data_path, "rb") as fp:
+      images = np.load(fp)
+    images = np.transpose(images, (1, 0, 2, 3))
+    images = np.expand_dims(images, axis=-1)
+    for sequence in images:
+      yield dict(image_sequence=sequence)