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()