code files

app.py

@@ -1,17 +1,15 @@
 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
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
 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 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 Surya-OCR model
+# Load the second OCR model (Surya-OCR)
 surya_processor = TrOCRProcessor.from_pretrained("suryavarmaaddala/suryaocr")
 surya_model = VisionEncoderDecoderModel.from_pretrained("suryavarmaaddala/suryaocr")
 
@@ -19,19 +17,17 @@ 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))
+        image = Image.fromarray(image).convert("RGB")
     return image
 
-def got_ocr(image):
+def microsoft_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)
+    pixel_values = ms_processor(image, return_tensors="pt").pixel_values
     
-    with torch.no_grad():
-        output = got_model(input_tensor)
+    generated_ids = ms_model.generate(pixel_values)
+    generated_text = ms_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
     
-    return output[0].item()
+    return generated_text
 
 def surya_ocr(image):
     image = preprocess_image(image)
@@ -57,10 +53,10 @@ def search_text(text, query):
 
 def process_and_search(image, search_query):
     try:
-        got_score = got_ocr(image)
+        ms_text = microsoft_ocr(image)
         surya_text = surya_ocr(image)
         
-        result = f"GOT OCR Score: {got_score:.4f}\n\nExtracted Text:\n{surya_text}"
+        result = f"Microsoft OCR Result:\n{ms_text}\n\nSurya OCR Result:\n{surya_text}"
         processed_text = post_process_text(result)
         
         search = None