Merge pull request #701 from ChanchalKumarMaji:read-from-archive

PiperOrigin-RevId: 255552399

Author: copybara-github

tensorflow_datasets/image/downsampled_imagenet.py

@@ -19,10 +19,6 @@
 from __future__ import division
 from __future__ import print_function
 
-import os
-
-import tensorflow as tf
-
 from tensorflow_datasets.core import api_utils
 import tensorflow_datasets.public_api as tfds
 
@@ -107,38 +103,32 @@ def _split_generators(self, dl_manager):
     train_url = _DL_URL + "train_" + self.builder_config.name + ".tar"
     valid_url = _DL_URL + "valid_" + self.builder_config.name + ".tar"
 
-    extracted_paths = dl_manager.download_and_extract({
-        "train_images": train_url,
-        "valid_images": valid_url,
-    })
+    train_path, valid_path = dl_manager.download([
+        train_url,
+        valid_url,
+    ])
 
     return [
         tfds.core.SplitGenerator(
             name=tfds.Split.TRAIN,
             num_shards=10,
             gen_kwargs={
-                "path":
-                    os.path.join(extracted_paths["train_images"],
-                                 "train_" + self.builder_config.name),
+                "archive": dl_manager.iter_archive(train_path),
             }),
         tfds.core.SplitGenerator(
             name=tfds.Split.VALIDATION,
             num_shards=1,
             gen_kwargs={
-                "path":
-                    os.path.join(extracted_paths["valid_images"],
-                                 "valid_" + self.builder_config.name),
+                "archive": dl_manager.iter_archive(valid_path),
             }),
     ]
 
-  def _generate_examples(self, path):
-    images = tf.io.gfile.listdir(path)
-
-    for image in images:
+  def _generate_examples(self, archive):
+    for fname, fobj in archive:
       record = {
-          "image": os.path.join(path, image),
+          "image": fobj,
       }
       if self.version.implements(tfds.core.Experiment.S3):
-        yield image, record
+        yield fname, record
       else:
         yield record

tensorflow_datasets/image/downsampled_imagenet_test.py

@@ -34,10 +34,10 @@ class DownsampledImagenetTest(testing.DatasetBuilderTestCase):
       tfds.Split.VALIDATION: 2,
   }
 
-  DL_EXTRACT_RESULT = {
-      "train_images": "train_images",
-      "valid_images": "valid_images",
-  }
+  DL_EXTRACT_RESULT = [
+      "train_32x32.tar",
+      "valid_32x32.tar",
+  ]
 
 
 class DownsampledImagenetS3Test(DownsampledImagenetTest):