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

update

Browse Source
This commit is contained in:
luni
2022-09-13 09:23:01 +00:00
parent 977e49900d
commit 440444cd82
19 changed files with 1630 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dataset.py
+51 -38
View File
@@ -54,18 +54,22 @@ class DatasetSimple(tf.data.Dataset):
class StaticDataset():
def __init__(self, folder, signature) -> None:
def __init__(self, folder, signature, use_filled=True) -> None:
self.store = dict()
self.signature = signature
self.folder = folder
self.possible_keys = ['uv_static', 'mesh_static', 'mesh_mask', 'uv_mask', 'uv_shaped_body']
self.keys = [value for value in signature.keys() if value in self.possible_keys]
filled = ''
if use_filled:
filled = '_filled'
def read(self, id):
data = dict()
if 'uv_static' in self.signature.keys():
uv_static = np.load(os.path.join(self.folder, 'uv_static', id + '.npz')); uv_static = uv_static['uv_static']; uv_static = np.transpose(uv_static, [2, 0, 1]); data['uv_static']=uv_static
uv_static = np.load(os.path.join(self.folder, 'uv_static' + filled, id + '.npz')); uv_static = uv_static['uv_static']; uv_static = np.transpose(uv_static, [2, 0, 1]); data['uv_static' + filled]=uv_static
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
if 'mesh_mask' in self.signature.keys():
@@ -127,9 +131,12 @@ class DatasetFactory():
return tf_signature
def get_dataset(self, folder, batch_size, signature, shuffle=True):
def get_dataset(self, folder, batch_size, signature, shuffle=True, use_filled=True):
# static_data = StaticDataset(folder, signature)
dataset_idx = DatasetIdx(folder, shuffle)
filled = ''
if use_filled:
filled = '_filled'
def read_item(frame):
frame_split = tf.strings.split(frame, sep='_')
@@ -161,7 +168,7 @@ class DatasetFactory():
# 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])
uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient' + filled, 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])
item['uv_body_posed'] = uv_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']
@@ -179,10 +186,14 @@ class DatasetFactory():
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_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])
uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape' + filled, 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])
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']
uv_curvature = np.load(os.path.join(folder, 'uv_curvature' + filled, id + '_' + frame_nr + '.npz'));
if filled != '':
uv_curvature = uv_curvature['uv_curvature' + filled]
else:
uv_curvature = uv_curvature['curvature_uv']
item['uv_curvature'] = uv_curvature.astype(np.int32)
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']
@@ -239,58 +250,60 @@ if __name__ == '__main__':
signature_dynamic = { 'id': (),
# 'pose_shape': (72 + 11,),
# 'uv_posed': (3, 512, 256,),
# 'uv_body_posed': (3, 512, 256,),
'uv_body_posed': (3, 512, 256,),
# 'mesh_body_posed': (14475, 3,),
# 'mesh_posed': (14475, 3,),
'mesh_unposed': (14475, 3,),
# 'mesh_unposed': (14475, 3,),
'uv_unposed': (3, 512, 256,),
# 'uv_curvature': (512, 256),
# 'mesh_curvature': (14475)
}
signature_static = { #'uv_static': (3, 512, 256,),
# 'mesh_static': (14475, 3,),
'mesh_mask': (14475,),
# 'mesh_mask': (14475,),
'uv_mask': (512, 256,),
# 'uv_shaped_body': (3, 512, 256,),
}
dataset = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train',
batch_size=1,
signature=signature_dynamic)
static_dataset = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train',
signature=signature_static)
dataset = factory.get_dataset(folder='/home/workspace/local-data/train_100',
batch_size=300,
signature=signature_dynamic,
use_filled=False)
static_dataset = StaticDataset(folder='/home/workspace/local-data/train_100',
signature=signature_static,
use_filled=False)
# sum_uv = np.zeros((3, 512, 256), dtype=np.float32)
# sum_mask = np.zeros((512, 256), dtype=np.float32)
# for i, batch in enumerate(dataset):
# print(i)
# batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
# sum_uv += batch['uv_unposed'].numpy().sum(axis=0)
# sum_mask += batch['uv_mask'].numpy().sum(axis=0)
sum_uv = np.zeros((3, 512, 256), dtype=np.float32)
sum_mask = np.zeros((512, 256), dtype=np.float32)
for i, batch in enumerate(dataset):
print(i)
batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
sum_uv += batch['uv_unposed'].numpy().sum(axis=0)
sum_mask += batch['uv_mask'].numpy().sum(axis=0)
# avg_jumpsuit = sum_uv / sum_mask[np.newaxis, :, :]
# sum_mask[sum_mask == 0] = 1
sum_mask[sum_mask == 0] = 1
avg_jumpsuit = sum_uv / sum_mask[np.newaxis, :, :]
# x = np.transpose(avg_jumpsuit, [1,2,0])
# x = uv2mesh(avg_jumpsuit[tf.newaxis,:,:,:])[0]
# go.Figure(data=[
# go.Mesh3d(
# # 8 vertices of a cube
# x=x[:,0],
# y=x[:,1],
# z=x[:,2],
# # i, j and k give the vertices of triangles
# i = F[:,0],
# j = F[:,1],
# k = F[:,2]
# )]).show()
x = np.transpose(avg_jumpsuit, [1,2,0])
x = uv2mesh(avg_jumpsuit[tf.newaxis,:,:,:])[0]
go.Figure(data=[
go.Mesh3d(
# 8 vertices of a cube
x=x[:,0],
y=x[:,1],
z=x[:,2],
# i, j and k give the vertices of triangles
i = F[:,0],
j = F[:,1],
k = F[:,2]
)]).show()
# np.save('custom_data/jumpsuit_20_avg_uv.npy', np.transpose(avg_jumpsuit, [1,2,0]))
np.save('custom_data/train_100_no_shape_avg_uv.npy', np.transpose(avg_jumpsuit, [1,2,0]))
# print()
print()