in memory storing, avg skirt calc
This commit is contained in:
+24
-19
@@ -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()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user