small updates

data/dataset_cape.py

@@ -29,7 +29,7 @@ def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file=
 
         if is_cape:
             frames = ['.'.join(frame.split('.')[:-1]) for frame in frames] # .split('_')[0]
-            frames = [Frame(frame.split('.')[0], frame.split('.')[1]) for frame in frames]
+            frames = [Frame(frame.split('.')[0], frame.split('.')[1], 'cape') for frame in frames]
         else:
             frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
             frames = [Frame(frame.split('_')[0], frame.split('_')[1]) for frame in frames]
@@ -91,11 +91,10 @@ class StaticCache():
 
         if 'mesh_static' in self.signature:
             mesh_static = np.load(os.path.join(self.folder, 'mesh_static', id +  '.npz')); mesh_static = mesh_static['mesh_static']; data['mesh_static'] = mesh_static.astype(np.float32)
-        if 'mesh_mask' in self.signature:
-            if self.is_cape:
-                data['mesh_mask'] = np.ones((3638), np.float32)
-            else:
-                mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32); data['mesh_mask'] = mesh_mask
+        if self.is_cape:
+            data['mesh_mask'] = np.ones((3638), np.float32)
+        else:
+            mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32); data['mesh_mask'] = mesh_mask
         # if 'uv_shaped_body' in self.signature:
         #     uv_shaped_body = np.load(os.path.join(self.folder, 'uv_shaped_body', id + '.npz')); uv_shaped_body = uv_shaped_body['uv_shaped_body']; data['uv_shaped_body'] = uv_shaped_body
 
@@ -198,7 +197,12 @@ class DatasetCape(torch.utils.data.Dataset):
                 item['uv_posed' + res_string] = uv_posed.astype(np.float32)
 
             if 'uv_normals' + res_string in self.signature:
-                uv_normals = np.load(os.path.join(self.folder, 'uv_normals' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_normals = uv_normals['uv_normals']
+                try:
+                    uv_normals = np.load(os.path.join(self.folder, 'uv_normals' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_normals = uv_normals['uv_normals']
+                except:
+                    print('Warning: uv_normals not found for frame ' + frame.id + separator + frame.frame_nr + '.npz')
+                    uv_normals = np.zeros((res, int(res/2), 3), dtype=np.float32)
+
                 item['uv_normals' + res_string] = uv_normals
 
 
@@ -223,27 +227,28 @@ if __name__ == '__main__':
 
     signature = { 'id', 
                   'frame_nr',
-                  'pose_shape',
-                  'uv_posed',
-                  'uv_body_posed',
-                  'mesh_body_posed',
-                  'mesh_posed',
-                  'mesh_unposed',
-                  'uv_unposed',
-                  'uv_curvature',
-                  'mesh_curvature',
-                  'uv_static',
-                  'mesh_static',
-                  'mesh_mask',
-                  'uv_mask',
-                  'uv_shaped_body',
+                #   'pose_shape',
+                #   'uv_posed',
+                #   'uv_body_posed',
+                #   'mesh_body_posed',
+                #   'mesh_posed',
+                #   'mesh_unposed',
+                #   'uv_unposed',
+                  'uv_normals',
+                #   'mesh_curvature',
+                #   'uv_static',
+                #   'mesh_static',
+                #   'mesh_mask',
+                #   'uv_mask',
+                #   'uv_shaped_body',
                   }
     
-    dataset = Dataset('/data-local/127/train', signature, resolutions=[512], device=device, frames_file='frames_jumpsuit.txt')
+    dataset = DatasetCape('/data-local/127_2/cape_processed/train/female', signature, resolutions=[512, 256, 128, 64], device=device, frames_file=None, is_cape=True)
     train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)
 
-    for batch in train_dataloader:
+    for i, batch in enumerate(train_dataloader):
         batch
+        print(i)
         pass
 
     pass 

data/utils.py

@@ -11,7 +11,9 @@ class Frame:
 
     def __eq__(self, other):
         return (self.id, self.frame_nr, self.garment) == (other.id, other.frame_nr, self.garment)
-
+    
+    def __str__(self):
+            return f'{self.id}_{self.frame_nr}_{self.garment}'
 
 def save_to_h5file(h5file, id, frame_nr, garment, data, name, type):
     filters = tb.Filters(complevel=9, complib='blosc2')