init

dataset.py

@@ -0,0 +1,41 @@
+import os
+import sys
+import time
+import cv2
+import numpy as np
+import tensorflow as tf
+
+class DatasetSimple(tf.data.Dataset):
+
+    def _generator(folder):
+        # Opening the file
+        folder = folder.decode("utf-8")
+        uv_folder = os.path.join(folder, 'npz') 
+        pose_hape_folder = os.path.join(folder, 'poses_shapes') 
+        frames = os.listdir(uv_folder)
+
+        for frame in frames:
+            uv = np.load(os.path.join(uv_folder, frame))
+            uv = uv['uv_gt']
+            uv = cv2.resize(uv, [64, 128])
+            uv = np.transpose(uv, [2, 0, 1])
+
+            pose_shape = np.load(os.path.join(pose_hape_folder, frame.split('.')[0] + '.npy'))
+
+            yield uv, pose_shape
+
+    def __new__(cls, folder, res=64):
+        return tf.data.Dataset.from_generator(
+            cls._generator,
+            output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (83), dtype = tf.float32)),
+            args=(folder, )
+        )
+
+if __name__ == '__main__':
+    
+    dataset = DatasetSimple('/workspace/Dataset/stylegan/npz_1')
+
+    for d in dataset:
+        d
+        pass
+