curvature training
This commit is contained in:
+2
-2
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user