This commit is contained in:
2022-10-04 21:50:58 +00:00
parent 5b82041acb
commit 3c798b473d
8 changed files with 270 additions and 125 deletions
+5 -3
View File
@@ -59,8 +59,10 @@ class StaticDataset():
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']
assert set(resolutions).issubset(set([32, 64, 128, 256, 512]))
self.possible_keys = ['uv_static', 'uv_static_32', 'uv_static_64', 'uv_static_128', 'uv_static_256',
'uv_mask', 'uv_mask_32', 'uv_mask_64', 'uv_mask_128', 'uv_mask_256',
'mesh_static', 'mesh_mask', 'uv_shaped_body']
self.keys = [value for value in signature.keys() if value in self.possible_keys]
if 'uv_static_128' in signature.keys() and 'uv_static_128' not in self.keys:
self.keys.append('uv_static_128')
@@ -146,7 +148,7 @@ class DatasetFactory():
def get_dataset(self, folder, batch_size, signature, shuffle=True, use_filled=False, resolutions=[512], in_memory=True):
# assert (use_filled and len(resolutions) >= 1) or (not use_filled and len(resolutions) == 1)
assert set(resolutions).issubset(set([32, 128, 512]))
assert set(resolutions).issubset(set([32, 64, 128, 256, 512]))
# static_data = StaticDataset(folder, signature)
dataset_idx = DatasetIdx(folder, shuffle)