Update app.py

app.py

@@ -11,13 +11,14 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)  # Automatyczne wykrycie t
 
 # Reguły do wykrywania NIP, kwot, dat
 nip_pattern = re.compile(r'\bPL\s?\d{10}\b|\b\d{10}\b')
-kwota_pattern = re.compile(r'\b\d+[\.,]?\d*\b')
-data_pattern = re.compile(r'\b\d{2}\.\d{2}\.\d{4}\b')
-payment_keywords = ["data płatności", "termin płatności", "zapłata", "płatność"]
+kwota_pattern = re.compile(r'\b\d+[\.,]?\d*\s?(PLN|zł|EUR|USD)?\b')  # Rozpoznawanie walut
+data_pattern = re.compile(r'\b\d{2}\.\d{2}\.\d{4}\b')  # Format DD.MM.YYYY
+payment_keywords = ["data płatności", "termin płatności", "zapłata", "zapłacono", "płatność"]
+seller_keywords = ["sprzedawca", "faktura wystawiona przez", "wystawca", "nazwa firmy"]
 
 def extract_invoice_data(pdf_file):
     with pdfplumber.open(pdf_file) as pdf:
-        words, boxes = [], []
+        words, boxes, full_text = [], [], []
 
         for page in pdf.pages:
             extracted_words = page.extract_words()
@@ -26,8 +27,14 @@ def extract_invoice_data(pdf_file):
                 bbox = [int(word['x0']), int(word['top']), int(word['x1']), int(word['bottom'])]  # Zaokrąglamy wartości
                 boxes.append(bbox)  # Pobieramy bounding box (pozycję słowa na stronie)
 
+            page_text = page.extract_text()
+            if page_text:
+                full_text.append(page_text.lower())
+
+    full_text = "\n".join(full_text)  # Łączymy cały tekst dokumentu
+
     # Tokenizacja tekstu + dodanie bounding boxes
-    encoding = tokenizer.encode_plus(words, boxes=boxes, return_tensors="pt", truncation=True)  # Poprawiona linia
+    encoding = tokenizer.encode_plus(words, boxes=boxes, return_tensors="pt", truncation=True)
 
     # Predykcja modelu
     with torch.no_grad():
@@ -40,22 +47,32 @@ def extract_invoice_data(pdf_file):
         if pred > 0:  # Pomijamy tło
             entities.append((token, model.config.id2label[pred]))
 
-    # Wyszukiwanie kluczowych wartości
+    # 🏢 Wyszukiwanie nazwy sprzedawcy
     seller_name = [token for token, label in entities if "ORG" in label]
-    seller_nip = nip_pattern.search(" ".join(words))
-    kwoty = kwota_pattern.findall(" ".join(words))
-    kwoty = [float(k.replace(",", ".")) for k in kwoty if k.replace(",", ".").replace(".", "").isdigit()]
-    total_amount = max(kwoty) if kwoty else None
+    
+    # Jeśli model nie znalazł, szukamy w tekście
+    if not seller_name:
+        for line in full_text.split("\n"):
+            if any(keyword in line for keyword in seller_keywords):
+                seller_name = line.split(":")[-1].strip()
+                break
+
+    # 🔢 Wyszukiwanie NIP
+    seller_nip = nip_pattern.search(full_text)
+    
+    # 💰 Wyszukiwanie kwoty całkowitej (największa kwota z walutą)
+    kwoty = kwota_pattern.findall(full_text)
+    kwoty = [k[0].replace(",", ".") for k in kwoty if k[0].replace(",", ".").replace(".", "").isdigit()]
+    total_amount = max(map(float, kwoty), default=None) if kwoty else None
 
-    # Szukamy daty płatności
+    # 📆 Wyszukiwanie daty płatności
     payment_date = None
-    for i, word in enumerate(words):
-        if any(keyword in word.lower() for keyword in payment_keywords):
-            if i + 1 < len(words):
-                date_match = data_pattern.search(words[i + 1])
-                if date_match:
-                    payment_date = date_match.group()
-                    break
+    for line in full_text.split("\n"):
+        if any(keyword in line for keyword in payment_keywords):
+            date_match = data_pattern.search(line)
+            if date_match:
+                payment_date = date_match.group()
+                break
 
     return {
         "Sprzedawca": " ".join(seller_name) if seller_name else "Nie znaleziono",