Add small images to parquet output.

PiperOrigin-RevId: 718878698
google-research · Jan 23, 2025 · 269b507 · 269b507
1 parent c575efe
commit 269b507
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/skai/generate_examples.py b/src/skai/generate_examples.py
@@ -1124,6 +1124,8 @@ def _example_to_dict(
     features['post_image_png_large'] = utils.get_bytes_feature(
         e, 'post_image_png_large'
     )[0]
+    features['pre_image_png'] = utils.get_bytes_feature(e, 'pre_image_png')[0]
+    features['post_image_png'] = utils.get_bytes_feature(e, 'post_image_png')[0]
   return features
 
 
@@ -1169,6 +1171,8 @@ def _write_examples_to_parquet(
       ('label', pyarrow.float64()),
       ('pre_image_png_large', pyarrow.binary()),
       ('post_image_png_large', pyarrow.binary()),
+      ('pre_image_png', pyarrow.binary()),
+      ('post_image_png', pyarrow.binary()),
   ])
   _ = (
       examples