import gradio as gr from transformers import AutoModel, AutoTokenizer import os import re # Load the OCR model and tokenizer tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True) model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, device_map='cuda', use_safetensors=True, pad_token_id=tokenizer.eos_token_id).eval().cuda() def extract_text_from_image(image): image_path = "temp_image.jpg" image.save(image_path) extracted_text = model.chat(tokenizer, image_path, ocr_type='ocr') os.remove(image_path) return extracted_text def search_and_highlight_keyword(extracted_text, keyword): if not keyword: return "
Please provide a keyword for searching.
" def highlight(match): # Custom background color and text color for highlighting return f"{match.group(0)}" pattern = re.compile(re.escape(keyword), re.IGNORECASE) highlighted_text = [] for line in extracted_text.splitlines(): if re.search(pattern, line): highlighted_line = re.sub(pattern, highlight, line) highlighted_text.append(highlighted_line) if highlighted_text: return 'Keyword '{keyword}' not found in the text.
" # Gradio interface components with gr.Blocks() as demo: gr.Markdown("# OCR and Keyword Search App with Highlighting") image_input = gr.Image(type="pil", label="Upload an Image (JPEG format)") text_output = gr.Textbox(label="Extracted Text", placeholder="Text will appear here after OCR.") extract_button = gr.Button("Extract Text") extract_button.click(fn=extract_text_from_image, inputs=image_input, outputs=text_output) keyword_input = gr.Textbox(label="Enter Keyword to Search and Highlight") search_result = gr.HTML(label="Highlighted Text with Keyword") search_button = gr.Button("Search and Highlight Keyword") search_button.click(fn=search_and_highlight_keyword, inputs=[text_output, keyword_input], outputs=search_result) demo.launch(share=True)