41 lines
1.1 KiB
Python
41 lines
1.1 KiB
Python
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
|
|
|