large refactor
This commit is contained in:
+42
-45
@@ -10,46 +10,46 @@ from custom_utils.custom_utils import nest_conditions, batch_to_dict
|
||||
|
||||
|
||||
|
||||
class DatasetSimple(tf.data.Dataset):
|
||||
# class DatasetSimple(tf.data.Dataset):
|
||||
|
||||
def _generator(folder):
|
||||
# Opening the file
|
||||
folder = folder.decode("utf-8")
|
||||
uv_folder = os.path.join(folder, 'uv_posed')
|
||||
pose_hape_folder = os.path.join(folder, 'poses_shapes')
|
||||
frames = os.listdir(uv_folder)
|
||||
random.shuffle(frames)
|
||||
# def _generator(folder):
|
||||
# # Opening the file
|
||||
# folder = folder.decode("utf-8")
|
||||
# uv_folder = os.path.join(folder, 'uv_posed')
|
||||
# pose_hape_folder = os.path.join(folder, 'poses_shapes')
|
||||
# frames = os.listdir(uv_folder)
|
||||
# random.shuffle(frames)
|
||||
|
||||
for frame in frames:
|
||||
id, frame_nr = frame.split('.')[0].split('_')
|
||||
uv_posed = np.load(os.path.join(folder, 'uv_posed', id + '_' + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']; uv_posed = np.transpose(uv_posed, [2, 0, 1])
|
||||
uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body']; uv_body_posed = np.transpose(uv_body_posed, [2, 0, 1])
|
||||
pose_shape = np.load(os.path.join(folder, 'pose_shape', id + '_' + frame_nr + '.npy'))
|
||||
uv_static = np.load(os.path.join(folder, 'uv_static', id + '.npz')); uv_static = uv_static['uv_static']; uv_static = np.transpose(uv_static, [2, 0, 1])
|
||||
mesh_static = np.load(os.path.join(folder, 'mesh_static', id + '.npz')); mesh_static = mesh_static['mesh_static']
|
||||
mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy'))
|
||||
mesh_mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32)
|
||||
uv_mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32)
|
||||
mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + '_' + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed']
|
||||
mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed']
|
||||
uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1])
|
||||
# for frame in frames:
|
||||
# id, frame_nr = frame.split('.')[0].split('_')
|
||||
# uv_posed = np.load(os.path.join(folder, 'uv_posed', id + '_' + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']; uv_posed = np.transpose(uv_posed, [2, 0, 1])
|
||||
# uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body']; uv_body_posed = np.transpose(uv_body_posed, [2, 0, 1])
|
||||
# pose_shape = np.load(os.path.join(folder, 'pose_shape', id + '_' + frame_nr + '.npy'))
|
||||
# uv_static = np.load(os.path.join(folder, 'uv_static', id + '.npz')); uv_static = uv_static['uv_static']; uv_static = np.transpose(uv_static, [2, 0, 1])
|
||||
# mesh_static = np.load(os.path.join(folder, 'mesh_static', id + '.npz')); mesh_static = mesh_static['mesh_static']
|
||||
# mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy'))
|
||||
# mesh_mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32)
|
||||
# uv_mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32)
|
||||
# mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + '_' + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed']
|
||||
# mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed']
|
||||
# uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1])
|
||||
|
||||
# uv_posed, (pose, shape, uv_static), mesh_static, mesh_body_posed, mask
|
||||
yield uv_posed, (pose_shape[:72], pose_shape[72:], uv_static, uv_mask, uv_body_posed), mesh_static, mesh_body_posed, mesh_mask, mesh_posed, mesh_unposed, uv_unposed
|
||||
# # uv_posed, (pose, shape, uv_static), mesh_static, mesh_body_posed, mask
|
||||
# yield uv_posed, (pose_shape[:72], pose_shape[72:], uv_static, uv_mask, uv_body_posed), mesh_static, mesh_body_posed, mesh_mask, mesh_posed, mesh_unposed, uv_unposed
|
||||
|
||||
def __new__(cls, folder, res=256):
|
||||
return tf.data.Dataset.from_generator(
|
||||
cls._generator,
|
||||
output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32),
|
||||
(tf.TensorSpec(shape = (72), dtype = tf.float32), tf.TensorSpec(shape = (11), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32) ),
|
||||
tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
|
||||
tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
|
||||
tf.TensorSpec(shape = (14475,), dtype = tf.float32),
|
||||
tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
|
||||
tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
|
||||
tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), ),
|
||||
args=(folder, )
|
||||
)
|
||||
# def __new__(cls, folder, res=256):
|
||||
# return tf.data.Dataset.from_generator(
|
||||
# cls._generator,
|
||||
# output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32),
|
||||
# (tf.TensorSpec(shape = (72), dtype = tf.float32), tf.TensorSpec(shape = (11), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32) ),
|
||||
# tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
|
||||
# tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
|
||||
# tf.TensorSpec(shape = (14475,), dtype = tf.float32),
|
||||
# tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
|
||||
# tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
|
||||
# tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), ),
|
||||
# args=(folder, )
|
||||
# )
|
||||
|
||||
|
||||
|
||||
@@ -59,7 +59,7 @@ class StaticDataset():
|
||||
self.signature = signature
|
||||
self.folder = folder
|
||||
assert set(resolutions).issubset(set([32, 128, 512]))
|
||||
self.possible_keys = ['uv_static', 'uv_static_32', 'uv_static_128' 'mesh_static', 'mesh_mask', 'uv_mask', 'uv_mask_32', 'uv_mask_128', 'uv_shaped_body']
|
||||
self.possible_keys = ['uv_static', 'uv_static_32', 'uv_static_128', 'mesh_static', 'mesh_mask', 'uv_mask', 'uv_mask_32', 'uv_mask_128', 'uv_shaped_body']
|
||||
self.keys = [value for value in signature.keys() if value in self.possible_keys]
|
||||
if 'uv_static_128' in signature.keys():
|
||||
self.keys.append('uv_static_128')
|
||||
@@ -78,7 +78,7 @@ class StaticDataset():
|
||||
res_string = '_' + str(res)
|
||||
|
||||
if 'uv_static' + res_string in self.signature.keys():
|
||||
uv_static = np.load(os.path.join(self.folder, 'uv_static' + self.filled + res_string, id + '.npz')); uv_static = uv_static['uv_static' + self.filled]; uv_static = np.transpose(uv_static, [2, 0, 1]); data['uv_static' + res_string] = uv_static
|
||||
uv_static = np.load(os.path.join(self.folder, 'uv_static' + self.filled + res_string, id + '.npz')); uv_static = uv_static['uv_static' + self.filled]; data['uv_static' + res_string] = uv_static
|
||||
if 'uv_mask' + res_string in self.signature.keys():
|
||||
uv_mask = np.load(os.path.join(self.folder, 'uv_mask' + res_string, id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); data['uv_mask' + res_string] = uv_mask
|
||||
|
||||
@@ -88,14 +88,11 @@ class StaticDataset():
|
||||
if 'mesh_mask' in self.signature.keys():
|
||||
mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32); data['mesh_mask'] = mesh_mask
|
||||
if 'uv_shaped_body' in self.signature.keys():
|
||||
uv_shaped_body = np.load(os.path.join(self.folder, 'uv_shaped_body', id + '.npz')); uv_shaped_body = uv_shaped_body['uv_shaped_body']; uv_shaped_body = np.transpose(uv_shaped_body, [2, 0, 1]); data['uv_shaped_body'] = uv_shaped_body
|
||||
uv_shaped_body = np.load(os.path.join(self.folder, 'uv_shaped_body', id + '.npz')); uv_shaped_body = uv_shaped_body['uv_shaped_body']; data['uv_shaped_body'] = uv_shaped_body
|
||||
|
||||
return data
|
||||
|
||||
def __getitem__(self, id) -> dict:
|
||||
# id = tf.strings.unicode_decode(id, 'UTF-8')
|
||||
# id = tf.strings.as_string(id)
|
||||
# id = id.numpy()
|
||||
if id not in self.store:
|
||||
self.store[id] = self.read(id)
|
||||
|
||||
@@ -170,7 +167,7 @@ class DatasetFactory():
|
||||
|
||||
item['pose_shape'] = pose_shape
|
||||
if 'uv_posed' in signature.keys():
|
||||
uv_posed = np.load(os.path.join(folder, 'uv_posed', id + '_' + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']; uv_posed = np.transpose(uv_posed, [2, 0, 1])
|
||||
uv_posed = np.load(os.path.join(folder, 'uv_posed', id + '_' + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']
|
||||
item['uv_posed'] = uv_posed
|
||||
# if 'uv_body_posed' in signature.keys():
|
||||
# uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body']; uv_body_posed = np.transpose(uv_body_posed, [2, 0, 1])
|
||||
@@ -218,11 +215,11 @@ class DatasetFactory():
|
||||
item['uv_curvature' + res_string] = uv_curvature.astype(np.int32)
|
||||
|
||||
if 'uv_unposed' + res_string in signature.keys():
|
||||
uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape' + filled + res_string, id + '_' + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape' + filled]; uv_unposed_no_shape = np.transpose(uv_unposed_no_shape, [2, 0, 1])
|
||||
uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape' + filled + res_string, id + '_' + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape' + filled]
|
||||
item['uv_unposed' + res_string] = uv_unposed_no_shape
|
||||
|
||||
if 'uv_body_posed' + res_string in signature.keys():
|
||||
uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient' + filled + res_string, id + '_' + frame_nr + '.npz')); uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient' + filled]; uv_posed_body_no_orient = np.transpose(uv_posed_body_no_orient, [2, 0, 1])
|
||||
uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient' + filled + res_string, id + '_' + frame_nr + '.npz')); uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient' + filled]
|
||||
item['uv_body_posed' + res_string] = uv_posed_body_no_orient
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user