Spaces:
Sleeping
Sleeping
File size: 3,177 Bytes
3bc9acc 10c178b 3bc9acc 10c178b 3bc9acc 10c178b 3bc9acc 10c178b 3bc9acc 10c178b 3bc9acc 10c178b 3bc9acc 10c178b 3bc9acc 10c178b 3bc9acc 10c178b 3bc9acc 10c178b 3bc9acc |
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 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
import gradio as gr
import numpy as np
from PIL import Image
import torch
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
import re
# Load the first OCR model (Microsoft's TrOCR)
ms_processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
ms_model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
# Load the second OCR model (Surya-OCR)
surya_processor = TrOCRProcessor.from_pretrained("suryavarmaaddala/suryaocr")
surya_model = VisionEncoderDecoderModel.from_pretrained("suryavarmaaddala/suryaocr")
def preprocess_image(image):
if isinstance(image, str):
image = Image.open(image).convert("RGB")
elif isinstance(image, np.ndarray):
image = Image.fromarray(image).convert("RGB")
return image
def microsoft_ocr(image):
image = preprocess_image(image)
pixel_values = ms_processor(image, return_tensors="pt").pixel_values
generated_ids = ms_model.generate(pixel_values)
generated_text = ms_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
return generated_text
def surya_ocr(image):
image = preprocess_image(image)
pixel_values = surya_processor(image, return_tensors="pt").pixel_values
generated_ids = surya_model.generate(pixel_values)
generated_text = surya_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
return generated_text
def post_process_text(text):
# Simple post-processing to split into lines
return '\n'.join(text.split('. '))
def search_text(text, query):
try:
pattern = re.compile(query, re.IGNORECASE)
lines = text.split('\n')
matching_lines = [line for line in lines if pattern.search(line)]
return '\n'.join(matching_lines) if matching_lines else "No matches found."
except re.error:
return "Invalid regex pattern. Please try again."
def process_and_search(image, search_query):
try:
ms_text = microsoft_ocr(image)
surya_text = surya_ocr(image)
result = f"Microsoft OCR Result:\n{ms_text}\n\nSurya OCR Result:\n{surya_text}"
processed_text = post_process_text(result)
search = None
if search_query:
search = search_text(processed_text, search_query)
return image, processed_text, search
except Exception as e:
return None, f"An error occurred: {str(e)}", None
with gr.Blocks() as demo:
with gr.Row():
with gr.Column(scale=1):
image_input = gr.Image(type="filepath", label="Upload your image")
search_query_input = gr.Textbox(label="Enter search query")
submit_button = gr.Button("Submit")
with gr.Column(scale=2):
displayed_image = gr.Image(label="Uploaded Image")
ocr_result = gr.Textbox(label="OCR Result", lines=10)
search_result = gr.Textbox(label="Search Result", lines=5)
submit_button.click(
fn=process_and_search,
inputs=[image_input, search_query_input],
outputs=[displayed_image, ocr_result, search_result]
)
if __name__ == "__main__":
demo.launch() |