static encoder runs, but doesn't work yet

data/dataset.py

@@ -37,6 +37,8 @@ def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file=
             frames = sorted(os.listdir(os.path.join(folder, 'mesh_static') ))
             frames = [frame.split('.')[0] + '_4' for frame in frames]
 
+        frames = [Frame(frame.split('_')[0], frame.split('_')[1]) for frame in frames]
+
     if frames_file is not None:
         frames = []
         with open(os.path.join(folder, frames_file), 'r') as f:
@@ -122,7 +124,7 @@ class StaticCache():
 
 
 class Dataset(Dataset):
-    def __init__(self, folder, signature, resolutions, frames_file='frames_all.txt', device='cuda', use_cache=False, is_cape=False):
+    def __init__(self, folder, signature, resolutions=[512], frames_file='frames_all.txt', device='cuda', use_cache=False, is_cape=False):
         self.folder = folder
         self.device = device
         self.is_cape = is_cape
@@ -210,6 +212,10 @@ class Dataset(Dataset):
 
         data = { key: (torch.from_numpy(value).to(self.device) if not (type(value) == str or type(value) == torch.Tensor) else value) for key, value in data.items()}
 
+        for key, value in data.items():
+            if key.startswith('uv_') and value.shape[-1] == 3:
+                data[key] = torch.permute(value, [2,0,1])
+
         if self.cache is not None:
             self.cache[frame] = data
         return data