code files

app.py

@@ -0,0 +1,92 @@
+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()

requirements.txt

@@ -0,0 +1,10 @@
+gradio
+Pillow
+surya-ocr
+torch
+transformers
+tiktoken
+torchvision
+verovio
+accelerate
+rapidfuzz