in memory storing, avg skirt calc

This commit is contained in:
2022-09-30 21:53:50 +00:00
parent dca5368704
commit 52e7275edb
+24 -19
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]) -> None:
def __init__(self, folder, signature, use_filled=False, resolutions=[512], in_memory=True) -> None:
self.store = dict()
self.signature = signature
self.folder = folder
self.in_memory = in_memory
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.keys = [value for value in signature.keys() if value in self.possible_keys]
@@ -94,7 +95,11 @@ class StaticDataset():
def __getitem__(self, id) -> dict:
if id not in self.store:
self.store[id] = self.read(id)
item = self.read(id)
if not self.in_memory:
return item
else:
self.store[id] = item
return self.store[id]
@@ -270,9 +275,9 @@ if __name__ == '__main__':
import plotly.graph_objects as go
# import plotly.express as ex
# import cmapy as cm
uv2mesh = UV2Mesh()
uv2mesh = UV2Mesh(is_skirt=True)
F = np.load("custom_utils/smpl_faces.npy")
F = np.load("custom_utils/smpl_faces_skirt.npy")
# dataset = DatasetIdx('/home/workspace/data/dataset/stylegan_new/train')
@@ -284,37 +289,37 @@ if __name__ == '__main__':
# 'pose_shape': (72 + 11,),
# 'uv_posed': (3, 512, 256,),
#'uv_body_posed': (3, 512, 256,),
'uv_body_posed_32': (3, 32, 16,),
# 'uv_body_posed_32': (3, 32, 16,),
# 'mesh_body_posed': (14475, 3,),
# 'mesh_posed': (14475, 3,),
# 'mesh_unposed': (14475, 3,),
#'uv_unposed': (3, 512, 256,),
'uv_unposed_32': (3, 32, 16,),
'uv_curvature': (512, 256, 7),
'uv_unposed': (512, 256, 3),
# 'uv_unposed_32': (3, 32, 16,),
# 'uv_curvature': (512, 256, 7),
# 'mesh_curvature': (14475)
}
signature_static = { #'uv_static': (3, 512, 256,),
'uv_static_32': (3, 32, 16,),
# 'uv_static_32': (3, 32, 16,),
# 'mesh_static': (14475, 3,),
# 'mesh_mask': (14475,),
'uv_mask': (512, 256,),
# 'uv_shaped_body': (3, 512, 256,),
}
dataset = factory.get_dataset(folder='/home/workspace/local-data/train_100',
batch_size=4,
dataset = factory.get_dataset(folder='/home/workspace/local-data/train_skirt_25_2',
batch_size=300,
signature=signature_dynamic,
use_filled=True,
resolutions=[32])
static_dataset = StaticDataset(folder='/home/workspace/local-data/train_100',
use_filled=False,
resolutions=[512])
static_dataset = StaticDataset(folder='/home/workspace/local-data/train_skirt_25_2',
signature=signature_static,
use_filled=True,
resolutions=[32])
use_filled=False,
resolutions=[512])
sum_uv = np.zeros((3, 512, 256), dtype=np.float32)
sum_uv = np.zeros((512, 256, 3), dtype=np.float32)
sum_mask = np.zeros((512, 256), dtype=np.float32)
for i, batch in enumerate(dataset):
print(i)
@@ -323,7 +328,7 @@ if __name__ == '__main__':
sum_mask += batch['uv_mask'].numpy().sum(axis=0)
sum_mask[sum_mask == 0] = 1
avg_jumpsuit = sum_uv / sum_mask[np.newaxis, :, :]
avg_jumpsuit = sum_uv / sum_mask[:, :, np.newaxis]
x = np.transpose(avg_jumpsuit, [1,2,0])
x = uv2mesh(avg_jumpsuit[tf.newaxis,:,:,:])[0]
@@ -339,7 +344,7 @@ if __name__ == '__main__':
k = F[:,2]
)]).show()
np.save('custom_data/train_100_no_shape_avg_uv.npy', np.transpose(avg_jumpsuit, [1,2,0]))
np.save('custom_data/train_skirt_25_2_no_shape_avg_uv.npy', np.transpose(avg_jumpsuit, [1,2,0]))
print()