Add DMLab dataset to TFDS.

PiperOrigin-RevId: 279786128

Author: akolesnikoff

docs/release_notes.md

@@ -13,3 +13,5 @@
     ([guide](https://github.com/tensorflow/datasets/tree/master/docs/decode.md)).
 *   Add `duke_ultrasound` dataset of ultrasound phantoms and invivo liver images
     from the [MimickNet paper](https://arxiv.org/abs/1908.05782)
+*   Add Dmlab dataset from the
+    [VTAB benchmark](https://arxiv.org/abs/1910.04867).

tensorflow_datasets/image/__init__.py

@@ -42,6 +42,7 @@
 from tensorflow_datasets.image.cycle_gan import CycleGAN
 from tensorflow_datasets.image.deep_weeds import DeepWeeds
 from tensorflow_datasets.image.diabetic_retinopathy_detection import DiabeticRetinopathyDetection
+from tensorflow_datasets.image.dmlab import Dmlab
 from tensorflow_datasets.image.downsampled_imagenet import DownsampledImagenet
 from tensorflow_datasets.image.dsprites import Dsprites
 from tensorflow_datasets.image.dtd import Dtd

tensorflow_datasets/image/dmlab.py

@@ -0,0 +1,141 @@
+# coding=utf-8
+# Copyright 2019 The TensorFlow Datasets Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Dmlab dataset."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import io
+
+import os
+from absl import logging
+import tensorflow as tf
+
+import tensorflow_datasets.public_api as tfds
+
+_URL = "https://storage.googleapis.com/akolesnikov-dmlab-tfds/dmlab.tar.gz"
+
+
+class Dmlab(tfds.core.GeneratorBasedBuilder):
+  """Dmlab dataset."""
+
+  VERSION = tfds.core.Version("1.0.0")
+
+  def _info(self):
+    return tfds.core.DatasetInfo(
+        builder=self,
+        description=(r"""
+        The Dmlab dataset contains frames observed by the agent acting in the
+        DeepMind Lab environment, which are annotated by the distance between
+        the agent and various objects present in the environment. The goal is to
+        is to evaluate the ability of a visual model to reason about distances
+        from the visual input in 3D environments. The Dmlab dataset consists of
+        360x480 color images in 6 classes. The classes are
+        {close, far, very far} x {positive reward, negative reward}
+        respectively."""),
+        features=tfds.features.FeaturesDict({
+            "image": tfds.features.Image(shape=(360, 480, 3),
+                                         encoding_format="jpeg"),
+            "filename": tfds.features.Text(),
+            "label": tfds.features.ClassLabel(num_classes=6),
+        }),
+        homepage="https://github.com/google-research/task_adaptation",
+        citation=r"""@article{zhai2019visual,
+        title={The Visual Task Adaptation Benchmark},
+        author={Xiaohua Zhai and Joan Puigcerver and Alexander Kolesnikov and
+               Pierre Ruyssen and Carlos Riquelme and Mario Lucic and
+               Josip Djolonga and Andre Susano Pinto and Maxim Neumann and
+               Alexey Dosovitskiy and Lucas Beyer and Olivier Bachem and
+               Michael Tschannen and Marcin Michalski and Olivier Bousquet and
+               Sylvain Gelly and Neil Houlsby},
+                              year={2019},
+                              eprint={1910.04867},
+                              archivePrefix={arXiv},
+                              primaryClass={cs.CV},
+                              url = {https://arxiv.org/abs/1910.04867}
+                          }""",
+        supervised_keys=("image", "label")
+    )
+
+  def _split_generators(self, dl_manager):
+    path = dl_manager.download_and_extract(_URL)
+
+    return [
+        tfds.core.SplitGenerator(
+            name=tfds.Split.TRAIN,
+            gen_kwargs={
+                "images_dir_path": path,
+                "split_name": "train",
+            }),
+        tfds.core.SplitGenerator(
+            name=tfds.Split.VALIDATION,
+            gen_kwargs={
+                "images_dir_path": path,
+                "split_name": "validation",
+            }),
+        tfds.core.SplitGenerator(
+            name=tfds.Split.TEST,
+            gen_kwargs={
+                "images_dir_path": path,
+                "split_name": "test",
+            }),
+    ]
+
+  def _parse_single_image(self, example_proto):
+    """Parses single video from the input tfrecords.
+
+    Args:
+      example_proto: tfExample proto with a single video.
+
+    Returns:
+      dict with all frames, positions and actions.
+    """
+
+    feature_map = {
+        "image": tf.io.FixedLenFeature(shape=[], dtype=tf.string),
+        "filename": tf.io.FixedLenFeature(shape=[], dtype=tf.string),
+        "label": tf.io.FixedLenFeature(shape=[], dtype=tf.int64),
+    }
+
+    parse_single = tf.io.parse_single_example(example_proto, feature_map)
+
+    return parse_single
+
+  def _generate_examples(self, images_dir_path, split_name):
+    path_glob = os.path.join(images_dir_path,
+                             "dmlab-{}.tfrecord*".format(split_name))
+    files = tf.io.gfile.glob(path_glob)
+
+    logging.info("Reading data from %s.", ",".join(files))
+    with tf.Graph().as_default():
+      ds = tf.data.TFRecordDataset(files)
+      ds = ds.map(
+          self._parse_single_image,
+          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:
+          while True:
+            result = sess.run(iterator)
+            yield result["filename"], {
+                "image": io.BytesIO(result["image"]),
+                "filename": result["filename"],
+                "label": result["label"],
+            }
+
+        except tf.errors.OutOfRangeError:
+          return

tensorflow_datasets/image/dmlab_test.py

@@ -0,0 +1,37 @@
+# coding=utf-8
+# Copyright 2019 The TensorFlow Datasets Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for DMlab dataset."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow_datasets import testing
+from tensorflow_datasets.image import dmlab
+
+
+class DmlabDatasetTest(testing.DatasetBuilderTestCase):
+  DATASET_CLASS = dmlab.Dmlab
+
+  SPLITS = {
+      "train": 2,
+      "test": 2,
+      "validation": 2,
+  }
+
+
+if __name__ == "__main__":
+  testing.test_main()

tensorflow_datasets/testing/test_data/fake_examples/dmlab/dmlab-test.tfrecord-00000-of-00001

@@ -0,0 +1 @@
+https://storage.googleapis.com/akolesnikov-dmlab-tfds/dmlab.tar.gz 3017022789 638b18fa69a5d61bbc310cd0b87ac603a39f41de0fdb07b6e77de274a24480a4