From 52e7275edb52cc40b1d2c636b1254ea31ea9135f Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Fri, 30 Sep 2022 21:53:50 +0000 Subject: [PATCH] in memory storing, avg skirt calc --- dataset.py | 43 ++++++++++++++++++++++++------------------- 1 file changed, 24 insertions(+), 19 deletions(-) diff --git a/dataset.py b/dataset.py index 0079fcb..b75aca6 100644 --- a/dataset.py +++ b/dataset.py @@ -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()