from pickle import APPEND
import gradio as gr
from transformers import LayoutLMForTokenClassification, LayoutLMTokenizer
from PIL import Image

model_name = "microsoft/layoutlm-base-uncased"
tokenizer = LayoutLMTokenizer.from_pretrained(model_name)
model = LayoutLMForTokenClassification.from_pretrained(model_name)


def image_to_text(image):
    image = Image.open(image.name)
    inputs = tokenizer(image, return_tensors="pt", padding=True, truncation=True)
    outputs = model(**inputs)
    predicted_ids = outputs.logits.argmax(-1)
    predicted_text = tokenizer.decode(predicted_ids[0])
    return predicted_text

app= gr.Interface(
    fn=image_to_text,
    inputs=gr.Image(label="Sélectionnez une image"),
    outputs=gr.Textbox(label="Texte extrait"),
    title="Image vers Texte ",
    description="Téléchargez une image et cliquez sur 'Soumettre' pour extraire le texte.",
)


app.launch()