Feat(predict_all): now returns result as dict

asheswook · Feb 19, 2023 · 6c9bf97 · 6c9bf97
1 parent 3679aa8
commit 6c9bf97
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/FaceFlow/main.py b/FaceFlow/main.py
@@ -280,8 +280,11 @@ def predict(self, image):
 
 
     def predict_all(self):
+        result = {}
         for filename in os.listdir(self.predict_dir):
-            if filename.endswith(".jpg") or filename.endswith(".png") or filename.endswith(".jpeg"):
+            if self._isImage(filename):
                 image = face_recognition.load_image_file(os.path.join(self.predict_dir, filename))
-                names = self.predict(image)
-                print(f"예측 결과- {filename}: {names}")
+                numberLabel = self.predict(image)
+                nameLabel = self.convert_labelType(numberLabel, To.NAME)
+                result[filename] = nameLabel
+        return result