Spaces:
Sleeping
Sleeping
File size: 2,745 Bytes
53d5710 211ccfb 53d5710 211ccfb f4f10a9 53d5710 142aefd f4f10a9 abd170d 53d5710 f4f10a9 53d5710 abd170d f4f10a9 53d5710 f4f10a9 53d5710 abd170d |
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 64 65 66 67 68 69 70 71 72 73 |
import gradio as gr
from transformers import VisionEncoderDecoderModel, TrOCRProcessor,AutoTokenizer,ViTFeatureExtractor
from PIL import Image
import torch
def preprocess_image(image):
# Resize while maintaining aspect ratio
target_size = (224, 224)
original_size = image.size
# Calculate the new size while maintaining aspect ratio
aspect_ratio = original_size[0] / original_size[1]
if aspect_ratio > 1: # Width is greater than height
new_width = target_size[0]
new_height = int(target_size[0] / aspect_ratio)
else: # Height is greater than width
new_height = target_size[1]
new_width = int(target_size[1] * aspect_ratio)
# Resize the image
resized_img = image.resize((new_width, new_height))
# Calculate padding values
padding_width = target_size[0] - new_width
padding_height = target_size[1] - new_height
# Apply padding to center the resized image
pad_left = padding_width // 2
pad_top = padding_height // 2
pad_image = Image.new('RGB', target_size, (255, 255, 255)) # White background
pad_image.paste(resized_img, (pad_left, pad_top))
return pad_image
# Load model directly
from transformers import AutoTokenizer, AutoModel,ViTFeatureExtractor,TrOCRProcessor,VisionEncoderDecoderModel
tokenizer = AutoTokenizer.from_pretrained("aayushpuri01/TrOCR-Devanagari")
model1 = VisionEncoderDecoderModel.from_pretrained("aayushpuri01/TrOCR-Devanagari")
feature_extractor1 = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k')
processor1 = TrOCRProcessor(feature_extractor=feature_extractor1, tokenizer=tokenizer)
# tokenizer = AutoTokenizer.from_pretrained("paudelanil/trocr-devanagari")
# model = VisionEncoderDecoderModel.from_pretrained("paudelanil/trocr-devanagari")
# feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k')
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model1.to(device)
def predict(image):
# Preprocess the image
image = Image.open(image).convert("RGB")
image = preprocess_image(image)
pixel_values = processor1(image, return_tensors="pt").pixel_values.to(device)
# Generate text from the image
generated_ids = model1.generate(pixel_values)
generated_text = processor1.batch_decode(generated_ids, skip_special_tokens=True)[0]
return generated_text
# Create the Gradio interface
interface = gr.Interface(
fn=predict,
inputs=gr.Image(type="filepath"),
outputs="text",
title="Devanagari OCR with TrOCR",
description="Upload an image with Devanagari script and get the text prediction using a pre-trained Vision-Text model."
)
# Launch the interface
interface.launch(share=True) |