Update app.py

app.py

@@ -1,60 +1,40 @@
-# -*- coding: utf-8 -*-
-"""
-Final WebApp using Gradio
-"""
+# app.py
 
-# Required Libraries
 import cv2
-import torch
 from pytesseract import pytesseract
 from transformers import AutoModel, AutoTokenizer
 import gradio as gr
-import tempfile
-import os
 
-# Check if GPU is available
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-
-# Load models for OCR
-tokenizer_eng = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
-model_eng = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True).to(device).eval()
-
-# Tesseract configuration for Hindi OCR
-pytesseract.tesseract_cmd = '/usr/bin/tesseract'
-tesseract_config = '--oem 3 --psm 6 -l hin'
-
-# OCR function for both English and Hindi
-def perform_ocr(img, language):
-    # Use a temporary file for the uploaded image
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_img:
-        img.save(temp_img.name)
-        img_path = temp_img.name
+# Model and Tesseract Configuration
+def load_models():
+    tokenizer_eng = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
+    model_eng = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True).eval().cuda()
+    
+    pytesseract.tesseract_cmd = '/usr/bin/tesseract'
+    tesseract_config = '--oem 3 --psm 6 -l hin'
+    
+    return tokenizer_eng, model_eng, tesseract_config
+
+# Perform OCR Function
+def perform_ocr(img, language, model_eng, tesseract_config):
+    img_path = "/tmp/uploaded_image.png"
+    img.save(img_path)
 
     res_eng = ""
     res_hin = ""
 
     if language in ["English", "Both"]:
-        # Ensure that inference is done on the correct device (GPU or CPU)
-        with torch.no_grad():
-            # Move inputs to the appropriate device
-            try:
-                res_eng = model_eng.chat(tokenizer_eng, img_path, ocr_type='ocr')
-            except Exception as e:
-                print(f"Error during English OCR: {e}")
-                res_eng = "Error during English OCR"
+        res_eng = model_eng.chat(tokenizer_eng, img_path, ocr_type='ocr')
 
     if language in ["Hindi", "Both"]:
         img_cv = cv2.imread(img_path)
         res_hin = pytesseract.image_to_string(img_cv, config=tesseract_config)
 
-    # Cleanup temporary file
-    os.remove(img_path)
-
     return res_eng, res_hin
 
 # Keyword Search Functionality
-def ocr_and_search(image, language, keyword):
-    english_text, hindi_text = perform_ocr(image, language)
+def ocr_and_search(image, language, keyword, model_eng, tesseract_config):
+    english_text, hindi_text = perform_ocr(image, language, model_eng, tesseract_config)
 
     extracted_english = f"Extracted English Text:\n{english_text}" if english_text else "No English text extracted."
     extracted_hindi = f"Extracted Hindi Text:\n{hindi_text}" if hindi_text else "No Hindi text extracted."
@@ -64,6 +44,7 @@ def ocr_and_search(image, language, keyword):
     if keyword:
         if language in ["English", "Both"] and keyword.lower() in english_text.lower():
             search_results.append(f"Keyword '{keyword}' found in English text.")
+
         if language in ["Hindi", "Both"] and keyword.lower() in hindi_text.lower():
             search_results.append(f"Keyword '{keyword}' found in Hindi text.")
 
@@ -72,17 +53,25 @@ def ocr_and_search(image, language, keyword):
     return extracted_english, extracted_hindi, search_output
 
 # Gradio Interface Setup
-with gr.Blocks() as app:
-    gr.Markdown("### OCR Application")
-    image_input = gr.Image(type="pil", label="Upload Image")
-    language_selection = gr.Radio(choices=["English", "Hindi", "Both"], label="Select Language")
-    keyword_input = gr.Textbox(placeholder="Enter keyword to search", label="Keyword Search")
-    output_english = gr.Textbox(label="Extracted English Text", interactive=False)
-    output_hindi = gr.Textbox(label="Extracted Hindi Text", interactive=False)
-    output_search = gr.Textbox(label="Search Results", interactive=False)
-
-    submit_button = gr.Button("Submit")
-    submit_button.click(fn=ocr_and_search, inputs=[image_input, language_selection, keyword_input], outputs=[output_english, output_hindi, output_search])
-
-# Launch the Gradio app
-app.launch()
+def create_interface(model_eng, tesseract_config):
+    with gr.Blocks() as app:
+        gr.Markdown("### OCR Application")
+        image_input = gr.Image(type="pil", label="Upload Image")
+        language_selection = gr.Radio(choices=["English", "Hindi", "Both"], label="Select Language")
+        keyword_input = gr.Textbox(placeholder="Enter keyword to search", label="Keyword Search")
+        output_english = gr.Textbox(label="Extracted English Text", interactive=False)
+        output_hindi = gr.Textbox(label="Extracted Hindi Text", interactive=False)
+        output_search = gr.Textbox(label="Search Results", interactive=False)
+
+        submit_button = gr.Button("Submit")
+        submit_button.click(fn=ocr_and_search, inputs=[image_input, language_selection, keyword_input], outputs=[output_english, output_hindi, output_search])
+
+    return app
+
+def main():
+    tokenizer_eng, model_eng, tesseract_config = load_models()
+    app = create_interface(model_eng, tesseract_config)
+    app.launch()
+
+if __name__ == "__main__":
+    main()