curvature training

This commit is contained in:
2023-04-08 00:43:08 +02:00
parent 52e5658dcc
commit cada175296
7 changed files with 309 additions and 13 deletions
+2 -2
View File
@@ -186,7 +186,7 @@ class Dataset(Dataset):
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).to(self.device)
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:
@@ -206,7 +206,7 @@ class Dataset(Dataset):
data = { key: (torch.from_numpy(value) if not (type(value) == str or type(value) == torch.Tensor) else value) for key, value in item.items()}
for key, value in data.items():
if key.startswith('uv_') and value.shape[-1] == 3:
if key.startswith('uv_') and (value.shape[-1] == 3 or value.shape[-1] == 7):
data[key] = torch.permute(value, [2,0,1])
if self.cache is not None: