import gradio as gr
import easyocr
import re

# Initialize the EasyOCR reader for Hindi and English
reader = easyocr.Reader(['hi', 'en'])

# Step 1: Function to extract text using EasyOCR for both Hindi and English
def extract_text(image):
    results = reader.readtext(image, detail=0)  # 'detail=0' to get only the text output
    extracted_text = " ".join(results)  # Join all text results into a single string
    return extracted_text

# Step 2: Function to search and highlight the keyword in the extracted text
def search_keyword(extracted_text, keyword):
    if keyword.lower() in extracted_text.lower():
        # Highlight the keyword in the text using re for case-insensitive replacement
        highlighted_text = re.sub(f"(?i)({re.escape(keyword)})", r"[\1]", extracted_text)
        return highlighted_text
    else:
        # Return the extracted text with a note if the keyword is not found
        return f"Keyword '{keyword}' not found. Here is the extracted text:\n\n{extracted_text}"

# Gradio Interface
with gr.Blocks() as demo:
    with gr.Row():
        image_input = gr.Image(type="pil", label="Upload Image")
        extract_button = gr.Button("Extract Text")
    
    extracted_text_output = gr.Textbox(label="Extracted Text", interactive=False)

    with gr.Row():
        keyword_input = gr.Textbox(label="Enter Keyword")
        search_button = gr.Button("Search Keyword")
    
    highlighted_output = gr.Textbox(label="Result", interactive=False)

    # Step 1: When the user clicks the extract button, extract the text
    extract_button.click(fn=extract_text, inputs=image_input, outputs=extracted_text_output)

    # Step 2: When the user clicks the search button, search for the keyword
    search_button.click(fn=search_keyword, inputs=[extracted_text_output, keyword_input], outputs=highlighted_output)

# Launch the web app with sharing enabled
demo.launch(share=True)