import os
import gradio as gr
from transformers import TrOCRProcessor, TrOCRForConditionalGeneration
from PIL import Image
import torch

# 🛡️ Configuration du proxy si nécessaire
os.environ["HTTP_PROXY"] = "http://meditelproxy.meditel.int:80"
os.environ["HTTPS_PROXY"] = "http://meditelproxy.meditel.int:80"

# 🔄 Chargement du modèle et du processor
model_name = "microsoft/trocr-base-handwritten"
model = TrOCRForConditionalGeneration.from_pretrained(model_name)
processor = TrOCRProcessor.from_pretrained(model_name)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.eval()

# 🧠 Fonction OCR
def ocr_from_image(image_file, ocr_type):
    if image_file is None:
        return "Veuillez importer une image."

    # Prétraitement de l'image
    image = Image.open(image_file.name).convert("RGB")
    pixel_values = processor(images=image, return_tensors="pt").pixel_values.to(device)

    # Génération de texte
    with torch.no_grad():
        generated_ids = model.generate(pixel_values)

    # Décodage du texte généré
    generated_text = processor.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
    return generated_text

# 🔘 Types d’OCR (juste pour l’interface ici)
ocr_types = ["ocr", "format"]

# 🎨 Interface Gradio
iface = gr.Interface(
    fn=ocr_from_image,
    inputs=[
        gr.File(label="Importer une image", file_types=[".jpg", ".jpeg", ".png"]),
        gr.Radio(ocr_types, label="Type d’OCR", value="ocr")
    ],
    outputs="text",
    title="🧠 OCR manuscrit avec TrOCR",
    description="Importez une image manuscrite pour extraire le texte avec le modèle Microsoft TrOCR."
)

# 🚀 Lancement
if __name__ == "__main__":
    iface.launch()