File size: 1,244 Bytes
81e172d 427b187 81e172d 427b187 81e172d 427b187 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 |
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() |