add low res uv mask option + no filled check

This commit is contained in:
2022-09-19 11:02:32 +00:00
parent 37df8468e2
commit 4da4df0d61
+4 -4
View File
@@ -59,7 +59,7 @@ class StaticDataset():
self.signature = signature self.signature = signature
self.folder = folder self.folder = folder
assert set(resolutions).issubset(set([32, 128, 512])) 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_shaped_body'] 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] self.keys = [value for value in signature.keys() if value in self.possible_keys]
if 'uv_static_128' in signature.keys(): if 'uv_static_128' in signature.keys():
self.keys.append('uv_static_128') self.keys.append('uv_static_128')
@@ -79,14 +79,14 @@ class StaticDataset():
if 'uv_static' + res_string in self.signature.keys(): if 'uv_static' + res_string in self.signature.keys():
uv_static = np.load(os.path.join(self.folder, 'uv_static' + self.filled + res_string, id + '.npz')); uv_static = uv_static['uv_static' + self.filled]; uv_static = np.transpose(uv_static, [2, 0, 1]); data['uv_static' + res_string] = uv_static uv_static = np.load(os.path.join(self.folder, 'uv_static' + self.filled + res_string, id + '.npz')); uv_static = uv_static['uv_static' + self.filled]; uv_static = np.transpose(uv_static, [2, 0, 1]); data['uv_static' + res_string] = uv_static
if 'uv_mask' in self.signature.keys():
uv_mask = np.load(os.path.join(self.folder, 'uv_mask' + res_string, id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); data['uv_mask' + res_string] = uv_mask
if 'mesh_static' in self.signature.keys(): 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 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(): if 'mesh_mask' in self.signature.keys():
mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32); data['mesh_mask'] = mesh_mask mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32); data['mesh_mask'] = mesh_mask
if 'uv_mask' in self.signature.keys():
uv_mask = np.load(os.path.join(self.folder, 'uv_mask', id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); data['uv_mask'] = uv_mask
if 'uv_shaped_body' in self.signature.keys(): if 'uv_shaped_body' in self.signature.keys():
uv_shaped_body = np.load(os.path.join(self.folder, 'uv_shaped_body', id + '.npz')); uv_shaped_body = uv_shaped_body['uv_shaped_body']; uv_shaped_body = np.transpose(uv_shaped_body, [2, 0, 1]); data['uv_shaped_body'] = uv_shaped_body uv_shaped_body = np.load(os.path.join(self.folder, 'uv_shaped_body', id + '.npz')); uv_shaped_body = uv_shaped_body['uv_shaped_body']; uv_shaped_body = np.transpose(uv_shaped_body, [2, 0, 1]); data['uv_shaped_body'] = uv_shaped_body
@@ -143,7 +143,7 @@ class DatasetFactory():
def get_dataset(self, folder, batch_size, signature, shuffle=True, use_filled=True, resolutions=[512]): def get_dataset(self, folder, batch_size, signature, shuffle=True, use_filled=True, resolutions=[512]):
assert (use_filled and len(resolutions) >= 1) or (not use_filled and len(resolutions) == 1) # 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, 128, 512]))
# static_data = StaticDataset(folder, signature) # static_data = StaticDataset(folder, signature)