Create app.py

app.py

@@ -0,0 +1,74 @@
+import gradio as gr
+from transformers import AutoModel, AutoTokenizer
+import os
+import re
+
+# Load the model and tokenizer from local path
+# Assuming your model and tokenizer are stored in '/content/model' directory in Colab
+model_path = 'pranavdaware/web_ocr'
+
+# Load the model and tokenizer from the local directory
+tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+model = AutoModel.from_pretrained(model_path, trust_remote_code=True, low_cpu_mem_usage=True, device_map='cuda', use_safetensors=True)
+model = model.eval().cuda()
+
+# Function to extract text using OCR
+def ocr_processing(image_file):
+    try:
+        # Perform OCR on the uploaded image
+        result = model.chat(tokenizer, image_file, ocr_type='ocr')
+        return result
+    except Exception as e:
+        return str(e)
+
+# Function to search for keywords in extracted text
+def search_keyword(ocr_text, keyword):
+    try:
+        # Use regex to search for the keyword and highlight matches
+        matches = re.findall(rf"({keyword})", ocr_text, re.IGNORECASE)
+        if matches:
+            highlighted_text = re.sub(rf"({keyword})", r'<mark>\1</mark>', ocr_text, flags=re.IGNORECASE)
+            return highlighted_text
+        else:
+            return f"No matches found for '{keyword}' in the extracted text."
+    except Exception as e:
+        return str(e)
+
+# Gradio interface
+def main():
+    # Gradio app layout
+    with gr.Blocks() as demo:
+        gr.Markdown("# OCR and Keyword Search Application")
+
+        with gr.Row():
+            with gr.Column():
+                image_input = gr.Image(type="filepath", label="Upload your image")
+                keyword_input = gr.Textbox(label="Enter keyword to search")
+                ocr_output = gr.Textbox(label="OCR Output")
+                search_output = gr.HTML(label="Search Results")
+
+                # Button for OCR processing
+                process_button = gr.Button("Process Image for OCR")
+
+                # Connect the OCR processing function to the button
+                process_button.click(
+                    fn=ocr_processing,
+                    inputs=image_input,
+                    outputs=ocr_output
+                )
+
+                # Button for keyword search
+                search_button = gr.Button("Search Keyword in OCR Text")
+
+                # Connect the search function to the button
+                search_button.click(
+                    fn=search_keyword,
+                    inputs=[ocr_output, keyword_input],
+                    outputs=search_output
+                )
+
+    # Launch the Gradio demo
+    demo.launch()
+
+if __name__ == "__main__":
+    main()