app.py

from transformers import TrOCRProcessor, VisionEncoderDecoderModel
from PIL import Image
import pdf2image
import gradio as gr
import base64
import io
import torch

# ✅ Używamy lżejszego modelu dla większej wydajności
MODEL_NAME = "microsoft/trocr-base-stage1"
processor = TrOCRProcessor.from_pretrained(MODEL_NAME)
model = VisionEncoderDecoderModel.from_pretrained(MODEL_NAME)

# ✅ Funkcja przetwarzania faktury (odbiera Base64 PDF)
def process_invoice(pdf_base64):
    try:
        # ✅ Dekodowanie Base64 na plik PDF
        pdf_bytes = base64.b64decode(pdf_base64)
        pdf_stream = io.BytesIO(pdf_bytes)

        # ✅ Konwersja PDF do obrazu
        images = pdf2image.convert_from_bytes(pdf_stream.read())

        # ✅ Przetwarzanie pierwszej strony faktury
        image = images[0].convert("RGB")
        pixel_values = processor(images=image, return_tensors="pt").pixel_values

        # ✅ Uruchomienie modelu AI
        generated_ids = model.generate(pixel_values)
        extracted_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]

        return extracted_text
    except Exception as e:
        return f"Błąd przetwarzania: {str(e)}"

# ✅ API Gradio (teraz akceptuje Base64 jako wejście)
interface = gr.Interface(
    fn=process_invoice,
    inputs=gr.Textbox(label="PDF w Base64"),
    outputs="text",
    title="AI Invoice Processor",
    description="Prześlij fakturę jako Base64, a AI wyodrębni dane tekstowe.",
)

# ✅ Uruchomienie aplikacji
interface.launch()