Update app.py

app.py

@@ -61,12 +61,18 @@ class FirebaseConnector:
         cred = credentials.Certificate(firebase_config)
         firebase_admin.initialize_app(cred)
         self.db = firestore.client()
+        self.collecttion_name = os.getenv("collection", "id-card")
 
     def add(self, ocr: str, key: str):
-        doc_ref = self.db.collection("id-card").document(key)
+        doc_ref = self.db.collection(self.collecttion_name).document(key)
         if doc_ref.get().exists:
             raise ValueError(f"Document with key {key} already exists.")
-        return self.db.collection("id-card").add({"ocr": ocr}, document_id=key)
+        return self.db.collection(self.collecttion_name).add({"ocr": ocr}, document_id=key)
+
+    def get_id_counted(self):
+        docs = self.db.collection(self.collecttion_name).stream()
+        count = sum(1 for _ in docs)
+        return count
 
 
 def search_element(words: list[str], target: str, thr: float = 60):
@@ -79,8 +85,11 @@ class VoterCardVerifier:
     def __init__(self):
         self.db = FirebaseConnector()
         self.ocr_engine = EasyOCRReader()
+        self.thr = int(os.getenv("thr", 75))
+        self.n_cards = 0
 
-    def verify_card(self, image: Union[str, Image.Image], card_number, thr: float = 60):
+    def verify_card(self, image: Union[str, Image.Image], card_number,):
+        self.n_cards = self.get_id_counted()
         # Convertir l'image pour Groq
         if isinstance(image, str):
             image = Image.open(image)
@@ -88,8 +97,8 @@ class VoterCardVerifier:
         # Préparer la requête pour le modèle et l'appeler
         ocr_data = self.ocr_engine.chat(image)
 
-        score1 = search_element(ocr_data, "camer", thr)
-        score2 = search_element(ocr_data, card_number, thr)
+        score1 = search_element(ocr_data, "camer", self.thr)
+        score2 = search_element(ocr_data, card_number, self.thr)
 
         ocr_data = score1 + score2
         logger.info(str(ocr_data))
@@ -100,25 +109,23 @@ class VoterCardVerifier:
             if is_valid:
                 # Sauvegarder dans Firebase
                 self.save_card_number(ocr_data, card_number)
-                return "Your ID Card have been recorded, thank you !!!"
+                self.n_cards += 1
+                return f"Your ID Card have been recorded, thank you !!! Num Card: {self.n_cards}"
             else:
-                return "Numéro de carte non trouvé sur l'image"
-        except:
-            return "ID Card already saved in the database"
+                return f"Card number not found, please improve the quality of image passed. Num Card: {self.n_cards}"
+        except ValueError as ex:
+            return f"ID Card already saved in the database. Num Card: {self.n_cards}"
 
     def save_card_number(self, ocr_data, card_number):
         return self.db.add(str(ocr_data), str(card_number))
 
     def get_id_counted(self):
-        docs = self.db.db.collection(collecttion_name).stream()
-        count = sum(1 for _ in docs)
-        return count
+        return self.db.get_id_counted()
 
 
 # Interface Gradio
 def create_interface():
     verifier = VoterCardVerifier()
-    # n_cards = verifier.get_id_counted()
     description = (
         "Card ID Counter is an application designed to count voter card numbers "
         "by analyzing uploaded images using OCR technology. It ensures accuracy by "
@@ -146,7 +153,6 @@ def create_interface():
 
 
 if __name__ == "__main__":
-
     # Créer et lancer l'interface
     demo = create_interface()
     demo.launch()