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

prepare for cape

Browse Source
This commit is contained in:
luni
2022-10-19 18:43:44 +00:00
parent 1f85598db5
commit 5881f330f6
5 changed files with 1479 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dataset.py
+43 -18
View File
@@ -54,7 +54,7 @@ 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) -> None:
def __init__(self, folder, signature, use_filled=False, resolutions=[512], in_memory=True, is_cape=False) -> None:
self.store = dict()
self.signature = signature
self.folder = folder
@@ -68,9 +68,20 @@ class StaticDataset():
self.keys.append('uv_static_128')
self.resolutions = resolutions
self.filled = ''
self.is_cape = is_cape
if use_filled:
self.filled = '_filled'
if is_cape:
self.uv_masks = dict()
uv_mask = np.load('custom_data/garment_uv_mask_full.npy').astype(np.int32)
self.uv_masks[512] = uv_mask.astype(np.float32)
for res in [32, 64, 128, 256]:
down = tf.image.resize(uv_mask[:,:,tf.newaxis], size=(res, int(res/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
down = np.squeeze(down.numpy().astype(np.float32))
self.uv_masks[res] = down
def read(self, id):
data = dict()
@@ -83,13 +94,19 @@ class StaticDataset():
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
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
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
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
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)
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 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
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
@@ -118,10 +135,14 @@ class StaticDataset():
class DatasetIdx(tf.data.Dataset):
def __new__(cls, folder, shuffle=True):
def __new__(cls, folder, shuffle=True, is_cape=False):
uv_static_folder = os.path.join(folder, 'pose_shape')
frames = sorted(os.listdir(uv_static_folder))
frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
if is_cape:
frames = ['.'.join(frame.split('.')[:-1]) for frame in frames] # .split('_')[0]
else:
frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
if shuffle:
random.shuffle(frames)
@@ -146,18 +167,22 @@ class DatasetFactory():
return tf_signature
def get_dataset(self, folder, batch_size, signature, shuffle=True, use_filled=False, resolutions=[512], in_memory=False, do_repeat=True):
def get_dataset(self, folder, batch_size, signature, shuffle=True, use_filled=False, resolutions=[512], in_memory=False, do_repeat=True, is_cape=False):
# assert (use_filled and len(resolutions) >= 1) or (not use_filled and len(resolutions) == 1)
assert set(resolutions).issubset(set([32, 64, 128, 256, 512]))
# static_data = StaticDataset(folder, signature)
dataset_idx = DatasetIdx(folder, shuffle)
dataset_idx = DatasetIdx(folder, shuffle, is_cape=is_cape)
filled = ''
if use_filled:
filled = '_filled'
separator = '_'
if is_cape:
separator = '.'
def read_item(frame):
frame_split = tf.strings.split(frame, sep='_')
frame_split = tf.strings.split(frame, sep=separator)
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]
@@ -166,7 +191,7 @@ class DatasetFactory():
if 'id' in signature.keys():
item['id'] = id
if 'pose_shape' in signature.keys():
pose_shape = np.load(os.path.join(folder, 'pose_shape', id + '_' + frame_nr + '.npy'))
pose_shape = np.load(os.path.join(folder, 'pose_shape', id + separator + frame_nr + '.npy'))
##############################
pose_shape[:3] = 0
@@ -174,7 +199,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.load(os.path.join(folder, 'uv_posed', id + separator + 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])
@@ -187,10 +212,10 @@ class DatasetFactory():
# item['mesh_posed_body_no_orient'] = mesh_posed_body_no_orient
if 'mesh_body_posed' 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']
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']
item['mesh_body_posed'] = mesh_posed_body_no_orient
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']
mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + separator + 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']
@@ -199,12 +224,12 @@ class DatasetFactory():
# 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_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']
mesh_unposed_no_shape = np.load(os.path.join(folder, 'mesh_unposed_no_shape', 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 + '_' + frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature']
mesh_curvature = np.load(os.path.join(folder, 'mesh_curvature', id + separator + frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature']
item['mesh_curvature'] = mesh_curvature.astype(np.int32)
@@ -214,7 +239,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' + filled + res_string, id + '_' + frame_nr + '.npz'));
uv_curvature = np.load(os.path.join(folder, 'uv_curvature' + filled + res_string, id + separator + frame_nr + '.npz'));
if filled != '' or res_string != '':
uv_curvature = uv_curvature['uv_curvature' + filled]
else:
@@ -223,11 +248,11 @@ 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' + 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.load(os.path.join(folder, 'uv_unposed_no_shape' + filled + res_string, id + separator + 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.load(os.path.join(folder, 'uv_posed_body_no_orient' + filled + res_string, id + separator + 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