Spaces:
Sleeping
Sleeping
File size: 2,044 Bytes
e07af32 6c44e27 1e93254 8d587d1 0e323a0 9f8dfc0 e07af32 a509ba3 0e323a0 8d587d1 0e323a0 8d587d1 0e323a0 8d587d1 e07af32 8d587d1 0e323a0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
from transformers import MllamaForConditionalGeneration, AutoProcessor
from PIL import Image
import torch
import gradio as gr
import spaces
# Initialize model and processor
ckpt = "unsloth/Llama-3.2-11B-Vision-Instruct"
model = MllamaForConditionalGeneration.from_pretrained(
ckpt,
torch_dtype=torch.bfloat16
).to("cuda")
processor = AutoProcessor.from_pretrained(ckpt)
@spaces.GPU
def extract_text(image):
# Convert image to RGB
image = Image.open(image).convert("RGB")
# Create message structure
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "Extract handwritten text from the image and output only the extracted text without any additional description or commentary in output"},
{"type": "image"}
]
}
]
# Process input
texts = processor.apply_chat_template(messages, add_generation_prompt=True)
inputs = processor(text=texts, images=[image], return_tensors="pt").to("cuda")
# Generate output
outputs = model.generate(**inputs, max_new_tokens=250)
result = processor.decode(outputs[0], skip_special_tokens=True)
print(result)
# Clean up the output to remove the prompt and assistant text
if "assistant" in result.lower():
result = result[result.lower().find("assistant") + len("assistant"):].strip()
# Remove any remaining conversation markers
result = result.replace("user", "").replace("Extract handwritten text from the image and output only the extracted text without any additional description or commentary in output", "").strip()
print(result)
return result
# Create Gradio interface
demo = gr.Interface(
fn=extract_text,
inputs=gr.Image(type="filepath", label="Upload Image"),
outputs=gr.Textbox(label="Extracted Text"),
title="Handwritten Text Extractor",
description="Upload an image containing handwritten text to extract its content.",
)
# Launch the app
demo.launch(debug=True) |