Update app.py

app.py

@@ -1,50 +1,67 @@
+import cv2
+from pytesseract import pytesseract
+from transformers import AutoModel, AutoTokenizer
 import gradio as gr
-import torch
-import pytesseract
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import cv2  # Ensure you have OpenCV installed
-
-# Load models and tokenizers
-tokenizer_eng = AutoTokenizer.from_pretrained("ucaslcl/GOT-OCR2_0", trust_remote_code=True)
-tokenizer_hin = AutoTokenizer.from_pretrained("ucaslcl/GOT-OCR2_0", trust_remote_code=True)
-model_eng = AutoModelForCausalLM.from_pretrained("ucaslcl/GOT-OCR2_0", trust_remote_code=True)
-model_hin = AutoModelForCausalLM.from_pretrained("ucaslcl/GOT-OCR2_0", trust_remote_code=True)
-
-def perform_ocr(image):
-    """Perform OCR on the image for both English and Hindi."""
-    # Set device to CPU or GPU
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    
-    model_eng.to(device)
-    model_hin.to(device)
-
-    # Convert the input image to an OpenCV format
-    img_cv = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)  # Convert to BGR for OpenCV
-    img_path = "temp_image.jpg"  # Temporary path to save the image
-    cv2.imwrite(img_path, img_cv)  # Save the image temporarily
-
-    # Use pytesseract for English OCR
-    english_text = pytesseract.image_to_string(img_cv)
-
-    # Use pytesseract for Hindi OCR
-    tesseract_config = '--psm 6'
-    hindi_text = pytesseract.image_to_string(img_cv, lang='hin', config=tesseract_config)
-
-    return english_text, hindi_text
-
-def ocr_and_search(image):
-    """Process the image and extract text in both languages."""
-    english_text, hindi_text = perform_ocr(image)
-    return english_text, hindi_text
-
-# Gradio interface
-iface = gr.Interface(
-    fn=ocr_and_search,
-    inputs=gr.inputs.Image(type="numpy"),  # Use numpy array for OpenCV
-    outputs=["text", "text"],
-    title="OCR for English and Hindi",
-    description="Upload an image to extract text in English and Hindi."
-)
-
-# Launch the interface
-iface.launch()
+
+# Model and Tesseract Configuration
+# 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)
+
+    res_eng = ""
+    res_hin = ""
+
+    if language in ["English", "Both"]:
+        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)
+
+    return res_eng, res_hin
+
+# Keyword Search Functionality
+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."
+
+    # Search for the keyword in the extracted text
+    search_results = []
+    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.")
+
+    search_output = "\n".join(search_results) if search_results else "No matches found."
+
+    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])
+
+# Application Launch
+if __name__ == "__main__":
+    app.launch()