thierryguyot67 commited on
Commit
81e172d
·
verified ·
1 Parent(s): 8c72a19

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +36 -0
app.py ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoProcessor, AutoModelForDocumentClassification
3
+ from PIL import Image
4
+ import torch
5
+
6
+ # Charger le modèle et le processeur
7
+ model_name = "impira/layoutlm-document-classifier"
8
+ processor = AutoProcessor.from_pretrained(model_name)
9
+ model = AutoModelForDocumentClassification.from_pretrained(model_name)
10
+
11
+ def classify_document(image):
12
+ # Préparer l'image pour le modèle
13
+ inputs = processor(images=image, return_tensors="pt")
14
+
15
+ # Obtenir la prédiction
16
+ with torch.no_grad():
17
+ logits = model(**inputs).logits
18
+
19
+ # Convertir les logits en probabilités
20
+ probabilities = torch.nn.functional.softmax(logits, dim=-1)
21
+ predicted_class = torch.argmax(probabilities, dim=-1).item()
22
+
23
+ # Retourner la classe prédite et les probabilités
24
+ return f"Classe prédite : {model.config.id2label[predicted_class]}"
25
+
26
+ # Interface Gradio
27
+ interface = gr.Interface(
28
+ fn=classify_document,
29
+ inputs=gr.Image(type="pil", label="Téléchargez votre document"),
30
+ outputs="text",
31
+ title="Classification de Documents avec LayoutLM",
32
+ description="Ce modèle classifie les documents en fonction de leur contenu et de leur mise en page."
33
+ )
34
+
35
+ # Lancer l'application
36
+ interface.launch()