import gradio as gr
from transformers import AutoProcessor, LayoutLMForSequenceClassification
from PIL import Image
import torch

# Charger le modèle et le processeur
model_name = "impira/layoutlm-document-classifier"
processor = AutoProcessor.from_pretrained(model_name)
model = LayoutLMForSequenceClassification.from_pretrained(model_name)

def classify_document(image):
    # Préparer l'image pour le modèle
    inputs = processor(images=image, return_tensors="pt")
    
    # Obtenir la prédiction
    with torch.no_grad():
        logits = model(**inputs).logits

    # Convertir les logits en probabilités
    probabilities = torch.nn.functional.softmax(logits, dim=-1)
    predicted_class = torch.argmax(probabilities, dim=-1).item()

    # Retourner la classe prédite et les probabilités
    return f"Classe prédite : {model.config.id2label[predicted_class]}"

# Interface Gradio
interface = gr.Interface(
    fn=classify_document,
    inputs=gr.Image(type="pil", label="Téléchargez votre document"),
    outputs="text",
    title="Classification de Documents avec LayoutLM",
    description="Ce modèle classifie les documents en fonction de leur contenu et de leur mise en page."
)

# Lancer l'application
interface.launch()