Fix: check path exists when save model

asheswook · Mar 6, 2023 · e6947f0 · e6947f0
1 parent 08592ae
commit e6947f0
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/VisageSnap/main.py b/VisageSnap/main.py
@@ -155,6 +155,9 @@ def _load_model(self) -> LabelPropagation:
             return self.model
 
     def _save_model(self) -> None:
+        if not os.path.exists(os.path.join(os.getcwd(), "model")):
+            os.mkdir(os.path.join(os.getcwd(), "model"))
+
         data = (self.model, self.faces)
         with open(self.model_dir, "wb") as f:
             pickle.dump(data, f)