|
import os |
|
import gradio as gr |
|
from transformers import TrOCRProcessor, TrOCRForConditionalGeneration |
|
from PIL import Image |
|
import torch |
|
|
|
|
|
os.environ["HTTP_PROXY"] = "http://meditelproxy.meditel.int:80" |
|
os.environ["HTTPS_PROXY"] = "http://meditelproxy.meditel.int:80" |
|
|
|
|
|
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() |
|
|
|
|
|
def ocr_from_image(image_file, ocr_type): |
|
if image_file is None: |
|
return "Veuillez importer une image." |
|
|
|
|
|
image = Image.open(image_file.name).convert("RGB") |
|
pixel_values = processor(images=image, return_tensors="pt").pixel_values.to(device) |
|
|
|
|
|
with torch.no_grad(): |
|
generated_ids = model.generate(pixel_values) |
|
|
|
|
|
generated_text = processor.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0] |
|
return generated_text |
|
|
|
|
|
ocr_types = ["ocr", "format"] |
|
|
|
|
|
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." |
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
iface.launch() |
|
|