Upload tf_dataset.py with huggingface_hub

tf_dataset.py

@@ -0,0 +1,23 @@
+import tensorflow as tf
+import os
+import numpy as np
+
+def load_image(image_file):
+    image = tf.io.read_file(image_file)
+    image = tf.image.decode_jpeg(image, channels=3)
+    image = tf.image.convert_image_dtype(image, tf.float32)
+    image = tf.image.resize(image, [256, 256])
+    image = (image * 2) - 1
+    return image
+
+def get_dataset(root_path, subset="train"):
+    path_a = os.path.join(root_path, f"{subset}A")
+    path_b = os.path.join(root_path, f"{subset}B")
+    
+    list_a = tf.data.Dataset.list_files(path_a + "/*.jpg")
+    list_b = tf.data.Dataset.list_files(path_b + "/*.jpg")
+    
+    ds_a = list_a.map(load_image, num_parallel_calls=tf.data.AUTOTUNE)
+    ds_b = list_b.map(load_image, num_parallel_calls=tf.data.AUTOTUNE)
+    
+    return tf.data.Dataset.zip((ds_a, ds_b))