import gradio as gr
import easyocr
import re
from PIL import Image
import numpy as np

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

# Function to extract text using EasyOCR for both Hindi and English with error handling
def extract_text(image):
    try:
        # Resize the image to speed up processing (optional)
        image = image.resize((800, 800))  # Resize to 800x800 pixels
        # Convert PIL image to NumPy array
        image_np = np.array(image)
        # Extract text using EasyOCR
        results = reader.readtext(image_np, detail=0)
        extracted_text = " ".join(results)
        return extracted_text
    except Exception as e:
        # Return the error message if an exception occurs
        return f"Error occurred: {str(e)}"

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

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

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