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