improve dattaset cache

data/dataset.py

@@ -53,15 +53,15 @@ def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file=
 
 
 class StaticCache():
-    def __init__(self, folder, signature, resolutions=[512], in_memory=True, is_cape=False) -> None:
+    def __init__(self, folder, signature, resolutions=[512], is_cape=False, device='cuda') -> None:
         self.store = dict()
         self.signature = signature
         self.folder = folder
-        self.in_memory = in_memory
         assert set(resolutions).issubset(set([32, 64, 128, 256, 512]))
         self.resolutions = resolutions
         self.filled = ''
         self.is_cape = is_cape
+        self.device = device
         
         if is_cape:
             self.uv_masks = dict()
@@ -105,22 +105,14 @@ class StaticCache():
     def __getitem__(self, id) -> dict:
         if id not in self.store:
             item = self.read(id)
-            if not self.in_memory:
-                return item
-            else:
-                self.store[id] = item
-
-        return self.store[id]
-
-
-    def expand_batch(self, batch):
-        batch_size = batch['id'].shape[0]
-        items = []
-        for i in range(batch_size):
-            items.append(self.__getitem__(batch['id'][i].numpy().decode('UTF-8')))
-        for key in items[0].keys():
-            batch[key] = torch.stack([item[key] for item in items], dim=0)
-        return batch
+            stored_data = { key: torch.from_numpy(value).to('cpu') for key, value in item.items()}
+            for key, value in stored_data.items():
+                if key.startswith('uv_') and value.shape[-1] == 3:
+                    stored_data[key] = torch.permute(value, [2,0,1])
+            self.store[id] = stored_data
+        
+        data = { key: value.to(self.device) for key, value in self.store[id].items()}
+        return data
 
 
 class Dataset(Dataset):
@@ -148,7 +140,9 @@ class Dataset(Dataset):
         frame : Frame = self.frames[idx]
         if self.cache is not None:
             if frame in self.cache.keys():
-                return self.cache[frame]
+                data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in self.cache[frame].items()}
+                data = { **item, **self.static_cache[frame.id] }
+                return data
 
         separator = '_'
         if self.is_cape:
@@ -208,9 +202,8 @@ class Dataset(Dataset):
                 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
 
-        data = { **item, **self.static_cache[frame.id] }
 
-        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()}
+        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:
@@ -218,6 +211,9 @@ class Dataset(Dataset):
 
         if self.cache is not None:
             self.cache[frame] = data
+        
+        data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in data.items()}
+        data = { **item, **self.static_cache[frame.id] }
         return data