Spaces:
Paused
Paused
| import gradio as gr | |
| from transformers import AutoModel, AutoTokenizer | |
| import os | |
| import re # Import regular expressions module | |
| # 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() | |
| # Define the function to process images and extract text | |
| def extract_text_from_image(image): | |
| # Save the uploaded image temporarily | |
| image_path = "temp_image.jpg" | |
| image.save(image_path) | |
| # Call the model to perform OCR | |
| extracted_text = model.chat(tokenizer, image_path, ocr_type='ocr') | |
| # Remove the temporary image file | |
| os.remove(image_path) | |
| return extracted_text | |
| # Function to search for the keyword in extracted text and highlight it | |
| def search_and_highlight_keyword(extracted_text, keyword): | |
| if not keyword: | |
| return "<p>Please provide a keyword for searching.</p>" | |
| # Case-insensitive search and replace keyword with <mark> tag for highlighting | |
| def highlight(match): | |
| # Custom background color and text color for highlighting | |
| return f"<mark style='background-color: #ffcc00; color: black;'>{match.group(0)}</mark>" | |
| # Use regular expression to find the keyword in a case-insensitive manner | |
| pattern = re.compile(re.escape(keyword), re.IGNORECASE) | |
| highlighted_text = [] | |
| for line in extracted_text.splitlines(): # Split text into lines | |
| if re.search(pattern, line): # If keyword is found in the line | |
| highlighted_line = re.sub(pattern, highlight, line) # Highlight keyword | |
| highlighted_text.append(highlighted_line) | |
| if highlighted_text: | |
| return '<br>'.join(highlighted_text) # Join the lines with HTML <br> for line breaks | |
| else: | |
| return f"<p>Keyword '{keyword}' not found in the text.</p>" | |
| # Gradio interface components | |
| with gr.Blocks() as demo: | |
| # Image upload and OCR | |
| 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 search and highlight | |
| 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) | |
| # Launch the Gradio app | |
| demo.launch(share=True) | |