pdf-extractor / app.py
kryman27's picture
Update app.py
bf3bfc2 verified
raw
history blame
2.15 kB
import gradio as gr
import pdfplumber
import re
from transformers import pipeline
# Model do rozpoznawania nazw organizacji i wartości numerycznych
extractor = pipeline("ner", model="dbmdz/bert-large-cased-finetuned-conll03-english", aggregation_strategy="simple")
# Reguły do wykrywania NIP, kwot, dat
nip_pattern = re.compile(r'\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')
def extract_invoice_data(pdf_file):
with pdfplumber.open(pdf_file) as pdf:
full_text = "\n".join(page.extract_text() for page in pdf.pages if page.extract_text())
# Szukamy danych w tekście
entities = extractor(full_text)
seller_name = []
seller_nip = None
items = []
total_amount = None
invoice_date = None
for entity in entities:
if "ORG" in entity["entity_group"]:
seller_name.append(entity["word"]) # Zbieramy nazwę sprzedawcy
# Znajdujemy wartości numeryczne dla NIP, kwot, dat
seller_nip = nip_pattern.search(full_text)
invoice_date = data_pattern.search(full_text)
# **Naprawiamy błąd przetwarzania liczb**
kwoty = kwota_pattern.findall(full_text)
kwoty = [float(k.replace(",", ".")) for k in kwoty if k.replace(",", ".").replace(".", "").isdigit()] # Zamiana przecinka na kropkę
total_amount = max(kwoty) if kwoty else None # Pobranie największej wartości jako całkowita kwota faktury
return {
"Sprzedawca": " ".join(seller_name) if seller_name else "Nie znaleziono",
"NIP": seller_nip.group() if seller_nip else "Nie znaleziono",
"Data faktury": invoice_date.group() if invoice_date else "Nie znaleziono",
"Kwota całkowita": total_amount if total_amount else "Nie znaleziono"
}
# Interfejs użytkownika w Hugging Face Spaces
iface = gr.Interface(
fn=extract_invoice_data,
inputs=gr.File(label="Wybierz plik PDF"),
outputs="json",
title="Ekstrakcja danych z faktury",
description="Prześlij plik PDF, a model zwróci dane sprzedawcy, NIP, kwotę i datę faktury."
)
if __name__ == "__main__":
iface.launch()