Spaces:
Sleeping
Sleeping
import gradio as gr | |
from transformers import TrOCRProcessor, VisionEncoderDecoderModel | |
import requests | |
from PIL import Image | |
# Dictionary of model names and their corresponding HuggingFace model IDs | |
MODEL_OPTIONS = { | |
"Microsoft Handwritten": "microsoft/trocr-base-handwritten", | |
"Medieval Base": "medieval-data/trocr-medieval-base", | |
"Medieval Latin Caroline": "medieval-data/trocr-medieval-latin-caroline", | |
"Medieval Castilian Hybrida": "medieval-data/trocr-medieval-castilian-hybrida", | |
"Medieval Humanistica": "medieval-data/trocr-medieval-humanistica", | |
"Medieval Textualis": "medieval-data/trocr-medieval-textualis", | |
"Medieval Cursiva": "medieval-data/trocr-medieval-cursiva", | |
"Medieval Semitextualis": "medieval-data/trocr-medieval-semitextualis", | |
"Medieval Praegothica": "medieval-data/trocr-medieval-praegothica", | |
"Medieval Semihybrida": "medieval-data/trocr-medieval-semihybrida", | |
"Medieval Print": "medieval-data/trocr-medieval-print" | |
} | |
# Load image examples | |
urls = [ | |
'https://huggingface.co./medieval-data/trocr-medieval-base/resolve/main/images/caroline-1.png' | |
] | |
for idx, url in enumerate(urls): | |
image = Image.open(requests.get(url, stream=True).raw) | |
image.save(f"image_{idx}.png") | |
def load_model(model_name): | |
model_id = MODEL_OPTIONS[model_name] | |
processor = TrOCRProcessor.from_pretrained(model_id) | |
model = VisionEncoderDecoderModel.from_pretrained(model_id) | |
return processor, model | |
def process_image(image, model_name): | |
processor, model = load_model(model_name) | |
# prepare image | |
pixel_values = processor(image, return_tensors="pt").pixel_values | |
# generate (no beam search) | |
generated_ids = model.generate(pixel_values) | |
# decode | |
generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0] | |
return generated_text | |
title = "Interactive demo: TrOCR Model Switcher" | |
description = "Demo for the Medieval TrOCR HTR Models." | |
iface = gr.Interface( | |
fn=process_image, | |
inputs=[ | |
gr.Image(type="pil"), | |
gr.Dropdown(choices=list(MODEL_OPTIONS.keys()), label="Select Model") | |
], | |
outputs=gr.Textbox(), | |
title=title, | |
description=description, | |
examples=[ | |
["image_0.png", "Medieval Latin Caroline"] | |
] | |
) | |
iface.launch(debug=True) |