save

custom_utils/custom_utils.py

@@ -292,7 +292,7 @@ def sparse_to_tensor(L):
     return tf.SparseTensor(idx, data, shape)
 
 
-def load_model(model, checkpoint, disable_assert=False, assert_consumed_on=False, iteration=''):
+def load_model(model, checkpoint, disable_assert=False, assert_consumed_on=False, iteration='', return_iteration=False):
     ckpt = tf.train.Checkpoint(model=model)
     manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1)
     if not disable_assert:
@@ -306,7 +306,14 @@ def load_model(model, checkpoint, disable_assert=False, assert_consumed_on=False
             ckpt_ind = [i for i in range(len(manager.checkpoints)) if manager.checkpoints[i].endswith(name_base + '-' + iteration)][0]
             ckpt.restore(manager.checkpoints[ckpt_ind]).assert_existing_objects_matched()#.assert_consumed()
         else:
-            ckpt.restore(manager.latest_checkpoint).assert_existing_objects_matched()#.assert_consumed()
+            if not return_iteration:
+                ckpt.restore(manager.latest_checkpoint).assert_existing_objects_matched()#.assert_consumed()
+            else:
+                ckpt.restore(manager.latest_checkpoint).assert_existing_objects_matched()
+                if return_iteration:
+                    iteration = manager.latest_checkpoint.split('-')[-1]
+                    return iteration
+    
 
 
 def get_summary_str(model):
@@ -425,3 +432,14 @@ class Accumulator():
 
     def reset(self):
         self.dict = dict()
+
+
+from stl import mesh as mesh_stl
+
+def save_stl(file, mesh, faces):
+    obj = mesh_stl.Mesh(np.zeros(faces.shape[0], dtype=mesh_stl.Mesh.dtype))
+    for i, f in enumerate(faces):
+        for j in range(3):
+            obj.vectors[i][j] = mesh[f[j],:]    
+
+    obj.save(file)