cape dataset fix

This commit is contained in:
2023-07-04 18:21:08 +02:00
parent 2a4077303d
commit ffaf1a9a49
+5 -13
View File
@@ -82,7 +82,7 @@ class StaticCache():
res_string = '_' + str(res)
if 'uv_static' + res_string in self.signature:
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]; 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]; data['uv_static' + res_string] = uv_static.astype(np.float32)
if 'uv_mask' + res_string in self.signature:
if self.is_cape:
data['uv_mask' + res_string] = self.uv_masks[res]
@@ -158,6 +158,7 @@ class DatasetCape(torch.utils.data.Dataset):
item['id'] = frame.id
if 'frame_nr' in self.signature:
item['frame_nr'] = frame.frame_nr
item['garment'] = 'cape'
if 'pose_shape' in self.signature:
pose_shape = np.load(os.path.join(self.folder, 'pose_shape', frame.id + separator + frame.frame_nr + '.npy')).astype(np.float32)
pose_shape[:3] = 0
@@ -183,27 +184,18 @@ class DatasetCape(torch.utils.data.Dataset):
if res != 512:
res_string = '_' + str(res)
if 'uv_curvature' + res_string in self.signature:
uv_curvature = np.load(os.path.join(self.folder, 'uv_curvature' + res_string, frame.id + separator + frame.frame_nr + '.npz'));
if res_string != '':
uv_curvature = uv_curvature['uv_curvature']
else:
uv_curvature = uv_curvature['curvature_uv']
uv_curvature = torch.nn.functional.one_hot(torch.tensor(uv_curvature, device=self.device, dtype=torch.long), 7).float().to(self.device)
item['uv_curvature' + res_string] = uv_curvature
if 'uv_unposed' + res_string in self.signature:
uv_unposed_no_shape = np.load(os.path.join(self.folder, 'uv_unposed_no_shape' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape']
item['uv_unposed' + res_string] = uv_unposed_no_shape
item['uv_unposed' + res_string] = uv_unposed_no_shape.astype(np.float32)
if 'uv_body_posed' + res_string in self.signature:
uv_posed_body_no_orient = np.load(os.path.join(self.folder, uv_posed_body_no_orient_folder_name + res_string, frame.id + separator + frame.frame_nr + '.npz'))
uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']
item['uv_body_posed' + res_string] = uv_posed_body_no_orient
item['uv_body_posed' + res_string] = uv_posed_body_no_orient.astype(np.float32)
if 'uv_posed' + res_string in self.signature:
uv_posed = np.load(os.path.join(self.folder, 'uv_posed_no_orient' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_posed = uv_posed['uv_posed_no_orient']
item['uv_posed' + res_string] = uv_posed
item['uv_posed' + res_string] = uv_posed.astype(np.float32)
if 'uv_normals' + res_string in self.signature:
uv_normals = np.load(os.path.join(self.folder, 'uv_normals' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_normals = uv_normals['uv_normals']