Enable loading ImageNet file_name field to support loading ImageNet c…

…aptions PiperOrigin-RevId: 710515726
google-research · Dec 29, 2024 · b564cc4 · b564cc4
1 parent 97d6ac5
commit b564cc4
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/scenic/dataset_lib/big_transfer/bit.py b/scenic/dataset_lib/big_transfer/bit.py
@@ -85,7 +85,16 @@ def pp_fn(x, how):
     pp = builder.get_preprocess_fn(how)
     example = pp(x)
     # to scenic format
-    return {'inputs': example['image'], 'label': example['labels']}
+    if dataset_configs.dataset == 'imagenet2012' and 'file_name' in example:
+      return {
+          'inputs': example['image'],
+          'label': example['labels'],
+          'file_name': example['file_name'],
+      }
+    return {
+        'inputs': example['image'],
+        'label': example['labels'],
+    }
 
   # E.g. for testing with TAP.
   shuffle_buffer_size = (1000 if num_shards == 1 else

diff --git a/scenic/dataset_lib/big_transfer/builder.py b/scenic/dataset_lib/big_transfer/builder.py
@@ -23,7 +23,7 @@
 
 TPU_SUPPORTED_DTYPES = [
     tf.bool, tf.int32, tf.int64, tf.bfloat16, tf.float32, tf.complex64,
-    tf.uint32
+    tf.uint32, tf.string,
 ]