Create app.py

app.py

@@ -0,0 +1,36 @@
+import streamlit as st
+from PIL import Image
+import torch
+from transformers import LayoutLMv2ForSequenceClassification, LayoutLMv2Processor
+
+# Chargement du modèle et du processeur
+model_ft = LayoutLMv2ForSequenceClassification.from_pretrained("Tornaid/LayoutLMv2_D3_Classifier")
+processor_ft = LayoutLMv2Processor.from_pretrained("Tornaid/LayoutLMv2_D3_Classifier")
+
+label2id = {
+    'budget': 0, 'form': 1, 'file_folder': 2, 'invoice': 3, 'email': 4, 
+    'handwritten': 5, 'id_pieces': 6, 'advertisement': 7, 'carte postale': 8, 
+    'scientific_publication': 9, 'news_article': 10, 'scientific_report': 11, 
+    'resume': 12, 'letter': 13, 'presentation': 14, 'questionnaire': 15, 
+    'memo': 16, 'paye': 17, 'specification': 18
+}
+
+id2label = {id: label for label, id in label2id.items()}
+
+def predict_image_classification(image):
+    """Effectue la prédiction de classification sur l'image donnée."""
+    inputs = processor_ft(image, return_tensors="pt", truncation=True, max_length=512)
+    outputs = model_ft(**inputs)
+    prediction_index = outputs.logits.argmax(-1).item()
+    return id2label[prediction_index]
+
+# Interface Streamlit
+st.title("Classification de documents")
+
+uploaded_file = st.file_uploader("Choisissez un fichier image", type=["jpg", "jpeg", "png"])
+if uploaded_file is not None:
+    image = Image.open(uploaded_file).convert("RGB")
+    st.image(image, caption="Image chargée", use_column_width=True)
+    if st.button("Classer"):
+        label_pred = predict_image_classification(image)
+        st.write(f"Prédiction: {label_pred}")