File size: 1,349 Bytes
81f5391
 
c4c29da
 
81f5391
c4c29da
 
 
81f5391
c4c29da
 
 
 
 
 
e298b15
 
c4c29da
 
 
c7a515f
c4c29da
 
 
f5c48ae
81f5391
c4c29da
6370c40
e298b15
c4c29da
 
 
 
6370c40
c4c29da
 
6370c40
c4c29da
 
81f5391
c4c29da
 
 
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
import gradio as gr
from PIL import Image
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
import torch

# Load the improved TrOCR model
processor = TrOCRProcessor.from_pretrained("microsoft/trocr-large-handwritten")
model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-large-handwritten")

# Ensure model is on appropriate device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# OCR function
def extract_text(image):
    if image is None:
        return "Please upload an image."
    
    image = image.convert("RGB")
    pixel_values = processor(images=image, return_tensors="pt").pixel_values.to(device)

    generated_ids = model.generate(pixel_values)
    text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    
    return text.strip()

# Build Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("![Logo](logo.png)")
    gr.Markdown("""
    # Evaluator! Upgrade Your Writing Skill
    """)
    
    with gr.Row():
        input_image = gr.Image(type="pil", label="Upload Handwritten Image")
        output_text = gr.Textbox(label="Extracted Text", lines=8)

    submit_btn = gr.Button("Extract Text")
    submit_btn.click(fn=extract_text, inputs=input_image, outputs=output_text)

# Run the app
if __name__ == "__main__":
    demo.launch()