update div2k main, test data

Author: WilliamHYZhang

tensorflow_datasets/image/div2k.py

@@ -64,20 +64,16 @@ def __init__(self, data, **kwargs):
     if data not in _DATA_OPTIONS:
       raise ValueError("data must be one of %s" % _DATA_OPTIONS)
 
-    name = kwargs.get("name")
-    if name is None:
-      name = data
+    name = kwargs.get("name", data)
     kwargs["name"] = name
 
-    description = kwargs.get("description")
-    if description is None:
-      description = "Uses %s data." % data
+    description = kwargs.get("description", "Uses %s data." % data)
     kwargs["description"] = description
 
     super(Div2kConfig, self).__init__(**kwargs)
     self.data = data
 
-  def download_urls(self):
+  def download_urls():
     """Returns train and validation download urls for this config."""
     urls = {
         "train_lr_url": _DL_URLS["train_"+self.data],
@@ -91,14 +87,15 @@ def _make_builder_configs():
   configs = []
   for data in _DATA_OPTIONS:
     configs.append(Div2kConfig(
-        version=tfds.core.Version("1.0.0"),
+        version=tfds.core.Version("2.0.0"),
         data=data))
   return configs
 
 class Div2k(tfds.core.GeneratorBasedBuilder):
   """DIV2K dataset: DIVerse 2K resolution high quality images"""
 
   BUILDER_CONFIGS = _make_builder_configs()
+  VERSION = tfds.core.Version("2.0.0")
 
   def _info(self):
     return tfds.core.DatasetInfo(
@@ -108,14 +105,15 @@ def _info(self):
             "lr": tfds.features.Image(),
             "hr": tfds.features.Image(),
         }),
+        #homepage=_DL_URL,
         citation=_CITATION,
     )
 
   def _split_generators(self, dl_manager):
     """Returns SplitGenerators."""
 
     extracted_paths = dl_manager.download_and_extract(
-        self.builder_config.download_urls())
+        self.builder_config.download_urls)
 
     return [
         tfds.core.SplitGenerator(
@@ -139,12 +137,12 @@ def _generate_examples(self, lr_path, hr_path):
     if not tf.io.gfile.listdir(hr_path)[0].endswith(".png"):
       hr_path = os.path.join(hr_path, tf.io.gfile.listdir(hr_path)[0])
 
-    for root, dirs, files in tf.io.gfile.walk(lr_path):
-      if len(files) == 0:
-        continue
-      for file in files:
-        yield root + file, {
-            "lr": os.path.join(root, file),
-            "hr": os.path.join(hr_path, re.search(r'\d{4}',
-                                                  str(file)).group(0) + ".png")
-        }
+    for root, _, files in tf.io.gfile.walk(lr_path):
+      if len(files):
+        for file in files:
+          yield root + file, {
+              "lr": os.path.join(root, file),
+              "hr": os.path.join(hr_path,
+                                 re.search(r'\d{4}',
+                                           str(file)).group(0) + ".png")
+          }

tensorflow_datasets/testing/div2k.py

@@ -0,0 +1,67 @@
+"""Generates DIV2K like files with random data for testing."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import zipfile
+
+from absl import app
+from absl import flags
+
+import tensorflow as tf
+
+from tensorflow_datasets.core.utils import py_utils
+from tensorflow_datasets.testing import fake_data_utils
+
+flags.DEFINE_string("tfds_dir", py_utils.tfds_dir(),
+  "Path to tensorflow_datasets directory")
+
+FLAGS = flags.FLAGS
+
+DATA = {
+  "DIV2K_train_HR": "0001.png",
+  "DIV2K_train_LR_bicubic_X2": "0001x2.png",
+  "DIV2K_train_LR_bicubic_X3": "0001x3.png",
+  "DIV2K_train_LR_bicubic_X4": "0001x4.png",
+  "DIV2K_train_LR_difficult": "0001x4d.png",
+  "DIV2K_train_LR_mild": "0001x4m.png",
+  "DIV2K_train_LR_unknown_X2": "0001x2.png",
+  "DIV2K_train_LR_unknown_X3": "0001x3.png",
+  "DIV2K_train_LR_unknown_X4": "0001x4.png",
+  "DIV2K_train_LR_wild": "0001x4w.png",
+  "DIV2K_train_LR_x8": "0001x8.png",
+  "DIV2K_valid_HR": "0002.png",
+  "DIV2K_valid_LR_bicubic_X2": "0002x2.png",
+  "DIV2K_valid_LR_bicubic_X3": "0002x3.png",
+  "DIV2K_valid_LR_bicubic_X4": "0002x4.png",
+  "DIV2K_valid_LR_difficult": "0002x4d.png",
+  "DIV2K_valid_LR_mild": "0002x4m.png",
+  "DIV2K_valid_LR_unknown_X2": "0002x2.png",
+  "DIV2K_valid_LR_unknown_X3": "0002x3.png",
+  "DIV2K_valid_LR_unknown_X4": "0002x4.png",
+  "DIV2K_valid_LR_wild": "0002x4w.png",
+  "DIV2K_valid_LR_x8": "0002x8.png",  
+}
+
+def _output_dir():
+  """Returns output directory."""
+  return os.path.join(FLAGS.tfds_dir, "testing", "test_data", "fake_examples",
+    "div2k")
+
+def _generate_image(fdir, fname):
+  dirname = os.path.join(_output_dir(), fdir)
+  if not os.path.exists(dirname):
+    os.makedirs(dirname)
+  tf.io.gfile.copy(
+      fake_data_utils.get_random_png(1, 1),
+      os.path.join(dirname, fname),
+      overwrite=True)
+
+def main(argv):
+  for fdir, fname in DATA.items():
+    _generate_image(fdir, fname)
+
+if __name__ == "__main__":
+  app.run(main)