Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import numpy as np | |
| from PIL import Image | |
| import cv2 | |
| from transformers import TrOCRProcessor, VisionEncoderDecoderModel | |
| from huggingface_hub import hf_hub_download | |
| import torch | |
| import re | |
| # Download and load the GOT OCR model | |
| got_model_path = hf_hub_download(repo_id="junyeopkim/got_2.0_torch_script", filename="got_2.0_tiny.torchscript") | |
| got_model = torch.jit.load(got_model_path) | |
| # Load the Surya-OCR model | |
| 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(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) | |
| return image | |
| def got_ocr(image): | |
| image = preprocess_image(image) | |
| image = image.resize((224, 224)) | |
| input_tensor = torch.from_numpy(np.array(image)).permute(2, 0, 1).float() / 255.0 | |
| input_tensor = input_tensor.unsqueeze(0) | |
| with torch.no_grad(): | |
| output = got_model(input_tensor) | |
| return output[0].item() | |
| 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: | |
| got_score = got_ocr(image) | |
| surya_text = surya_ocr(image) | |
| result = f"GOT OCR Score: {got_score:.4f}\n\nExtracted Text:\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() |