trying to get basic model working

data/dataset.py

@@ -24,7 +24,7 @@ from cloth3d.util import quads2tris
 
 
 class Dataset(Dataset):
-    def __init__(self, folder, folder_rgb, frames_file='frames_all.txt', device='cuda'):
+    def __init__(self, folder, folder_rgb, frames_file='frames_all.txt', device='cuda', use_cache=True):
         self.folder = folder
         self.folder_rgb = folder_rgb
         self.device = device
@@ -37,6 +37,8 @@ class Dataset(Dataset):
         self.reader = DataReader(folder, 'superSMPL')
         
 
+        self.cache = dict() if use_cache else None
+
     def __len__(self):
         return len(self.frames)
 
@@ -44,6 +46,10 @@ class Dataset(Dataset):
     # TODO: make it work with two part outfits 
     def __getitem__(self, idx):
         frame : Frame = self.frames[idx]
+        if self.cache is not None:
+            if frame in self.cache.keys():
+                return self.cache[frame]
+
         path = os.path.join(self.folder, frame.id)
         info = self.reader.read_info(frame.id)
         # mesh_posed = readPC2Frame(os.path.join(path, frame.garment + '_registered.pc16'), frame.frame_nr, True)
@@ -189,7 +195,7 @@ class Dataset(Dataset):
 
 
         # image = self.transforms(image)
-        return { 'mesh_posed': mesh_posed,
+        data = { 'mesh_posed': mesh_posed,
                  'depth_image': depth_image, 
                  'mesh_mask': mesh_mask, 
                  'depth_mask': depth_mask, 
@@ -199,6 +205,9 @@ class Dataset(Dataset):
                  'rgb_image': rgb_image,
                 }
 
+        self.cache[frame] = data
+        return self.cache[frame]
+
 
 if __name__ == '__main__':
     from torch.utils.data import DataLoader