This commit is contained in:
2022-09-02 14:15:54 +00:00
parent ee8f1cb68c
commit 977e49900d
16 changed files with 2241 additions and 272 deletions
+19 -9
View File
@@ -99,11 +99,12 @@ class StaticDataset():
class DatasetIdx(tf.data.Dataset):
def __new__(cls, folder):
def __new__(cls, folder, shuffle=True):
uv_static_folder = os.path.join(folder, 'pose_shape')
frames = sorted(os.listdir(uv_static_folder))#[453:503]
frames = sorted(os.listdir(uv_static_folder))
frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
random.shuffle(frames)
if shuffle:
random.shuffle(frames)
return tf.data.Dataset.from_tensor_slices(frames)
@@ -126,9 +127,9 @@ class DatasetFactory():
return tf_signature
def get_dataset(self, folder, batch_size, signature):
def get_dataset(self, folder, batch_size, signature, shuffle=True):
# static_data = StaticDataset(folder, signature)
dataset_idx = DatasetIdx(folder)
dataset_idx = DatasetIdx(folder, shuffle)
def read_item(frame):
frame_split = tf.strings.split(frame, sep='_')
@@ -156,6 +157,9 @@ class DatasetFactory():
# if 'mesh_body_posed' in signature.keys():
# mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy'))
# item['mesh_body_posed'] = mesh_body_posed
# if 'mesh_posed_body_no_orient' in signature.keys():
# mesh_posed_body_no_orient = np.load(os.path.join(folder, 'mesh_posed_body_no_orient', id + '_' + frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient']
# item['mesh_posed_body_no_orient'] = mesh_posed_body_no_orient
if 'uv_body_posed' in signature.keys():
uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient', id + '_' + frame_nr + '.npz')); uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']; uv_posed_body_no_orient = np.transpose(uv_posed_body_no_orient, [2, 0, 1])
item['uv_body_posed'] = uv_posed_body_no_orient
@@ -165,12 +169,18 @@ class DatasetFactory():
if 'mesh_posed' in signature.keys():
mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + '_' + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed']
item['mesh_posed'] = mesh_posed
# if 'mesh_unposed' in signature.keys():
# mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed']
# item['mesh_unposed'] = mesh_unposed
# if 'uv_unposed' in signature.keys():
# 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])
# item['uv_unposed'] = uv_unposed
if 'mesh_unposed' in signature.keys():
mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed']
item['mesh_unposed'] = mesh_unposed
mesh_unposed_no_shape = np.load(os.path.join(folder, 'mesh_unposed_no_shape', id + '_' + frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape']
item['mesh_unposed'] = mesh_unposed_no_shape
if 'uv_unposed' in signature.keys():
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])
item['uv_unposed'] = uv_unposed
uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape', id + '_' + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape']; uv_unposed_no_shape = np.transpose(uv_unposed_no_shape, [2, 0, 1])
item['uv_unposed'] = uv_unposed_no_shape
if 'uv_curvature' in signature.keys():
uv_curvature = np.load(os.path.join(folder, 'uv_curvature', id + '_' + frame_nr + '.npz')); uv_curvature = uv_curvature['curvature_uv']
item['uv_curvature'] = uv_curvature.astype(np.int32)