save all

utils/postprocess.py

@@ -35,6 +35,8 @@ def fix_collision(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr=.004
         F_fix_coll_local = F_fix_coll_orig
     else:
         raise ValueError("V must have 14475 or 3638 vertices")
+    if V.shape[0] == 1:
+        F_fix_coll_local = F_fix_coll_local[:1]
     mesh = Meshes(B, F_fix_coll_local)
     B_vn = torch.stack(mesh.verts_normals_list())
    

utils/utils.py

@@ -1,9 +1,11 @@
 import os
 import sys
+import pickle
 import numpy as np
 from torch.utils.data import DataLoader
 from stl import mesh as mesh_stl
-
+from pathlib import Path
+from data.utils import Frame
 
 
 class Accumulator():
@@ -59,4 +61,29 @@ def save_stl(file, mesh, faces):
         for j in range(3):
             obj.vectors[i][j] = mesh[f[j],:]    
 
-    obj.save(file)
+    obj.save(file)
+
+
+def create_temp_frames(folder: str, frame: Frame):
+    with open(Path(folder) / 'frames_temp.txt', 'w') as f:
+        f.writelines([f'{frame}\n' for _ in range(4)])
+
+
+def pickle_load(file):
+    """
+    Load a pickle file.
+    """
+    with open(file, 'rb') as f:
+        loadout = pickle.load(f)
+
+    return loadout
+
+
+def pickle_dump(loadout, file):
+    """
+    Dump a pickle file. Create the directory if it does not exist.
+    """
+    os.makedirs(os.path.dirname(str(file)), exist_ok=True)
+
+    with open(file, 'wb') as f:
+        pickle.dump(loadout, f)