luni/cvpr
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

save

Browse Source
This commit is contained in:
luni
2023-03-10 17:25:44 +00:00
parent cb76d36c4f
commit 4b1cb0b06d
27 changed files with 1890 additions and 207 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dataset.py
+39 -22
View File
@@ -54,10 +54,11 @@ from custom_utils.custom_utils import nest_conditions, batch_to_dict
class StaticDataset():
def __init__(self, folder, signature, use_filled=False, resolutions=[512], in_memory=True, is_cape=False) -> None:
def __init__(self, folder, signature, folder_skirt='', use_filled=False, resolutions=[512], in_memory=True, is_cape=False) -> None:
self.store = dict()
self.signature = signature
self.folder = folder
self.folder_skirt = folder_skirt
self.in_memory = in_memory
assert set(resolutions).issubset(set([32, 64, 128, 256, 512]))
self.possible_keys = ['uv_static', 'uv_static_32', 'uv_static_64', 'uv_static_128', 'uv_static_256',
@@ -85,6 +86,9 @@ class StaticDataset():
def read(self, id):
data = dict()
temp_folder = self.folder
if not os.path.exists(os.path.join(temp_folder, 'uv_mask', id + '.npz')):
temp_folder = self.folder_skirt
for res in self.resolutions:
res_string = ''
@@ -92,23 +96,23 @@ 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]; data['uv_static' + res_string] = uv_static
uv_static = np.load(os.path.join(temp_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():
if self.is_cape:
data['uv_mask' + res_string] = self.uv_masks[res]
else:
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
uv_mask = np.load(os.path.join(temp_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
if 'mesh_static' in self.signature.keys():
mesh_static = np.load(os.path.join(self.folder, 'mesh_static', id + '.npz')); mesh_static = mesh_static['mesh_static']; data['mesh_static'] = mesh_static.astype(np.float32)
mesh_static = np.load(os.path.join(temp_folder, 'mesh_static', id + '.npz')); mesh_static = mesh_static['mesh_static']; data['mesh_static'] = mesh_static.astype(np.float32)
if 'mesh_mask' in self.signature.keys():
if self.is_cape:
data['mesh_mask'] = np.ones((3638), np.float32)
else:
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
mesh_mask = np.load(os.path.join(temp_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']; data['uv_shaped_body'] = uv_shaped_body
uv_shaped_body = np.load(os.path.join(temp_folder, 'uv_shaped_body', id + '.npz')); uv_shaped_body = uv_shaped_body['uv_shaped_body']; data['uv_shaped_body'] = uv_shaped_body
return data
@@ -135,7 +139,7 @@ class StaticDataset():
class DatasetIdx(tf.data.Dataset):
def __new__(cls, folder, shuffle=True, is_cape=False, testing=False, frames_file=None):
def __new__(cls, folder, shuffle=True, is_cape=False, testing=False, frames_file=None, folder_skirt=None, frames_file_skirt=''):
if frames_file is None:
uv_static_folder = os.path.join(folder, 'pose_shape')
frames = sorted(os.listdir(uv_static_folder))
@@ -156,6 +160,12 @@ class DatasetIdx(tf.data.Dataset):
lines = [line.strip() for line in lines]
frames = lines
if frames_file_skirt != '':
with open(os.path.join(folder_skirt, frames_file_skirt), 'r') as f:
lines = f.readlines()
lines = [line.strip() for line in lines]
frames.extend(lines)
if shuffle:
random.shuffle(frames)
@@ -194,10 +204,12 @@ class DatasetFactory():
use_one_hot=False,
use_no_orient=True,
testing=False,
frames_file=None):
frames_file=None,
frames_file_skirt='',
folder_skirt='',):
assert set(resolutions).issubset(set([32, 64, 128, 256, 512]))
dataset_idx = DatasetIdx(folder, shuffle, is_cape=is_cape, testing=testing, frames_file=frames_file)
dataset_idx = DatasetIdx(folder, shuffle, is_cape=is_cape, testing=testing, folder_skirt=folder_skirt, frames_file=frames_file, frames_file_skirt=frames_file_skirt)
separator = '_'
if is_cape:
@@ -222,6 +234,11 @@ class DatasetFactory():
id, frame_nr = frame_split[0].numpy().decode("utf-8"), frame_split[1].numpy().decode("utf-8")
# id, frame_nr = frame_split[0].numpy(), frame_split[1].numpy()
# static_info = static_data[id]
temp_folder = folder
if not os.path.exists(os.path.join(temp_folder, 'pose_shape', id + separator + frame_nr + '.npy')):
temp_folder = folder_skirt
item = dict()
if 'id' in signature.keys():
@@ -229,23 +246,23 @@ class DatasetFactory():
if 'frame_nr' in signature.keys():
item['frame_nr'] = frame_nr
if 'pose_shape' in signature.keys():
pose_shape = np.load(os.path.join(folder, 'pose_shape', id + separator + frame_nr + '.npy'))
pose_shape = np.load(os.path.join(temp_folder, 'pose_shape', id + separator + frame_nr + '.npy'))
if use_no_orient:
pose_shape[:3] = 0
item['pose_shape'] = pose_shape
if 'mesh_body_posed' in signature.keys():
if use_no_orient:
mesh_posed_body_no_orient = np.load(os.path.join(folder, 'mesh_posed_body_no_orient', id + separator + frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient']
mesh_posed_body_no_orient = np.load(os.path.join(temp_folder, 'mesh_posed_body_no_orient', id + separator + frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient']
item['mesh_body_posed'] = mesh_posed_body_no_orient
else:
mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy'))
mesh_body_posed = np.load(os.path.join(temp_folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy'))
item['mesh_body_posed'] = mesh_body_posed
if 'mesh_posed' in signature.keys():
if use_no_orient:
mesh_posed = np.load(os.path.join(folder, 'mesh_posed_no_orient', id + separator + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed_no_orient'].astype(np.float32)
mesh_posed = np.load(os.path.join(temp_folder, 'mesh_posed_no_orient', id + separator + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed_no_orient'].astype(np.float32)
item['mesh_posed'] = mesh_posed
else:
mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + separator + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed'].astype(np.float32)
mesh_posed = np.load(os.path.join(temp_folder, 'mesh_posed', id + separator + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed'].astype(np.float32)
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']
@@ -255,12 +272,12 @@ class DatasetFactory():
# item['uv_unposed'] = uv_unposed
if 'mesh_unposed' in signature.keys():
# print('one_hot')
mesh_unposed_no_shape = np.load(os.path.join(folder, 'mesh_unposed_no_shape' + one_hot, id + separator + frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape']
mesh_unposed_no_shape = np.load(os.path.join(temp_folder, 'mesh_unposed_no_shape' + one_hot, id + separator + frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape']
item['mesh_unposed'] = mesh_unposed_no_shape
if 'mesh_curvature' in signature.keys():
mesh_curvature = np.load(os.path.join(folder, 'mesh_curvature', id + separator + frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature']
mesh_curvature = np.load(os.path.join(temp_folder, 'mesh_curvature', id + separator + frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature']
item['mesh_curvature'] = mesh_curvature.astype(np.int32)
@@ -270,7 +287,7 @@ class DatasetFactory():
res_string = '_' + str(res)
if 'uv_curvature' + res_string in signature.keys():
uv_curvature = np.load(os.path.join(folder, 'uv_curvature' + res_string, id + separator + frame_nr + '.npz'));
uv_curvature = np.load(os.path.join(temp_folder, 'uv_curvature' + res_string, id + separator + frame_nr + '.npz'));
if res_string != '':
uv_curvature = uv_curvature['uv_curvature']
else:
@@ -279,24 +296,24 @@ class DatasetFactory():
item['uv_curvature' + res_string] = uv_curvature
if 'uv_unposed' + res_string in signature.keys():
uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape' + one_hot_2 + res_string, id + separator + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape']
uv_unposed_no_shape = np.load(os.path.join(temp_folder, 'uv_unposed_no_shape' + one_hot_2 + res_string, id + separator + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape']
item['uv_unposed' + res_string] = uv_unposed_no_shape
if 'uv_body_posed' + res_string in signature.keys():
if use_no_orient:
uv_posed_body_no_orient = np.load(os.path.join(folder, uv_posed_body_no_orient_folder_name + res_string, id + separator + frame_nr + '.npz'))
uv_posed_body_no_orient = np.load(os.path.join(temp_folder, uv_posed_body_no_orient_folder_name + res_string, id + separator + frame_nr + '.npz'))
uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']
item['uv_body_posed' + res_string] = uv_posed_body_no_orient
else:
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.load(os.path.join(temp_folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body']
item['uv_body_posed' + res_string] = uv_body_posed
if 'uv_posed' + res_string in signature.keys():
if use_no_orient:
uv_posed = np.load(os.path.join(folder, 'uv_posed_no_orient' + res_string, id + separator + frame_nr + '.npz')); uv_posed = uv_posed['uv_posed_no_orient']
uv_posed = np.load(os.path.join(temp_folder, 'uv_posed_no_orient' + res_string, id + separator + frame_nr + '.npz')); uv_posed = uv_posed['uv_posed_no_orient']
item['uv_posed' + res_string] = uv_posed
else:
uv_posed = np.load(os.path.join(folder, 'uv_posed' + res_string, id + separator + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']
uv_posed = np.load(os.path.join(temp_folder, 'uv_posed' + res_string, id + separator + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']
item['uv_posed' + res_string] = uv_posed
# print('here')