Update app.py

app.py

@@ -1,22 +1,19 @@
-# app.py
-
 import cv2
 from pytesseract import pytesseract
 from transformers import AutoModel, AutoTokenizer
 import gradio as gr
 
 # 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):
+# Load GOT2 model for English text and configure Tesseract for Hindi text
+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, device_map='cpu').eval()
+
+# Define Tesseract path and configuration for Hindi
+pytesseract.tesseract_cmd = '/usr/bin/tesseract'
+tesseract_config = '--oem 3 --psm 6 -l hin'
+
+# Perform OCR function
+def perform_ocr(img, language):
     img_path = "/tmp/uploaded_image.png"
     img.save(img_path)
 
@@ -33,8 +30,8 @@ def perform_ocr(img, language, model_eng, tesseract_config):
     return res_eng, res_hin
 
 # Keyword Search Functionality
-def ocr_and_search(image, language, keyword, model_eng, tesseract_config):
-    english_text, hindi_text = perform_ocr(image, language, model_eng, tesseract_config)
+def ocr_and_search(image, language, keyword):
+    english_text, hindi_text = perform_ocr(image, language)
 
     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."
@@ -53,25 +50,18 @@ def ocr_and_search(image, language, keyword, model_eng, tesseract_config):
     return extracted_english, extracted_hindi, search_output
 
 # Gradio Interface Setup
-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()
-
+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])
+
+# Application Launch
 if __name__ == "__main__":
-    main()
+    app.launch()