File size: 3,813 Bytes
34aedae |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
import gradio as gr
from gradio_client import Client, handle_file
import re
from thefuzz import fuzz
# hugging face clients for both OCR options
surya_ocr_client = Client("artificialguybr/Surya-OCR")
got_ocr_client = Client("stepfun-ai/GOT_official_online_demo")
# Global variable to store the extracted OCR text
extracted_text = ""
def ocr_extraction(image, ocr_model):
global extracted_text
if image is None:
return "Please upload an image first."
try:
if ocr_model == "Surya OCR":
client = surya_ocr_client
result = client.predict(
image=handle_file(image),
langs="en",
api_name="/ocr_workflow"
)
text_matches = re.findall(r"text='(.*?)'", str(result))
extracted_text = "\n".join(text_matches)
elif ocr_model == "GOT OCR":
client = got_ocr_client
result = client.predict(
image=handle_file(image),
got_mode="plain texts OCR",
fine_grained_mode="box",
ocr_color="red",
ocr_box="Hello!!",
api_name="/run_GOT"
)
extracted_text = result[0]
else:
return "Invalid OCR model selected."
return extracted_text
except Exception as e:
return f"An error occurred: {str(e)}"
def search_keyword(keyword, search_type):
global extracted_text
if not extracted_text:
return "No OCR text found. Please extract text from an image first."
if not keyword:
return extracted_text
if search_type == "Direct Search":
highlighted_text = re.sub(f"({re.escape(keyword)})", r'<span style="background-color: yellow;">\1</span>', extracted_text, flags=re.IGNORECASE)
else: # Nearest Search
words = extracted_text.split()
highlighted_words = []
for word in words:
if fuzz.ratio(word.lower(), keyword.lower()) >= 80: # Adjust threshold as needed
highlighted_words.append(f'<span style="background-color: yellow;">{word}</span>')
else:
highlighted_words.append(word)
highlighted_text = " ".join(highlighted_words)
return highlighted_text
with gr.Blocks(theme=gr.themes.Soft()) as gr_interface:
gr.Markdown("# 📷 OCR Text Extraction and Advanced Keyword Search 🔍")
with gr.Row():
with gr.Column(scale=1):
image_input = gr.Image(type="filepath", label="Upload Image")
ocr_model_dropdown = gr.Dropdown(
choices=["Surya OCR", "GOT OCR"],
value="Surya OCR",
label="Select OCR Model"
)
ocr_button = gr.Button("Extract Text", variant="primary")
with gr.Column(scale=2):
extracted_text_output = gr.Textbox(
label="Extracted Text",
placeholder="Text extracted from the image will appear here.",
lines=10
)
with gr.Row():
with gr.Column(scale=1):
keyword_input = gr.Textbox(label="Enter keyword to search")
search_type = gr.Radio(["Direct Search", "Nearest Search"], label="Search Type", value="Direct Search")
search_button = gr.Button("Search Keyword", variant="secondary")
with gr.Column(scale=2):
highlighted_output = gr.HTML(label="Highlighted Text")
ocr_button.click(
fn=ocr_extraction,
inputs=[image_input, ocr_model_dropdown],
outputs=extracted_text_output
)
search_button.click(
fn=search_keyword,
inputs=[keyword_input, search_type],
outputs=highlighted_output
)
gr_interface.launch(share=True) |