Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from transformers import AutoProcessor, Qwen2VLForConditionalGeneration | |
| from PIL import Image | |
| import torch | |
| import re | |
| # Load the pre-trained model and processor | |
| model = Qwen2VLForConditionalGeneration.from_pretrained( | |
| "Qwen/Qwen2-VL-2B-Instruct", | |
| torch_dtype="auto", | |
| device_map="auto", | |
| ) | |
| processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct") | |
| # Function to extract text from the image | |
| def extract_text(image): | |
| messages = [ | |
| { | |
| "role": "user", | |
| "content": [ | |
| {"type": "image"}, | |
| {"type": "text", "text": "can u extract the text in hindi"} | |
| ] | |
| } | |
| ] | |
| # Process input image and text prompt | |
| text_prompt = processor.apply_chat_template(messages, add_generation_prompt=True) | |
| inputs = processor( | |
| text=[text_prompt], | |
| images=[image], | |
| padding=True, | |
| return_tensors="pt" | |
| ) | |
| inputs = inputs.to("cuda" if torch.cuda.is_available() else "cpu") | |
| # Generate output text from the model | |
| output_ids = model.generate(**inputs, max_new_tokens=1024) | |
| generated_ids = [ | |
| output_ids[len(input_ids):] | |
| for input_ids, output_ids in zip(inputs.input_ids, output_ids) | |
| ] | |
| # Decode the generated text | |
| extracted_text = processor.batch_decode( | |
| generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True | |
| )[0] # Extracted text | |
| return extracted_text | |
| # Function to highlight keywords in the text, even for right-to-left scripts like Hindi | |
| def highlight_keywords(extracted_text, keywords): | |
| highlighted_text = extracted_text | |
| if keywords: | |
| for keyword in keywords.split(","): | |
| keyword = keyword.strip() | |
| if keyword: | |
| # Ensure correct Unicode support for keywords (use re.UNICODE for non-ASCII) | |
| highlighted_text = re.sub( | |
| re.escape(keyword), # Use re.escape to handle special characters in keywords | |
| r'<mark>\g<0></mark>', # Highlight the found keyword | |
| highlighted_text, | |
| flags=re.IGNORECASE | re.UNICODE # Ignore case, and handle Unicode characters | |
| ) | |
| return highlighted_text | |
| # First step: Extract text from the uploaded image | |
| def extract_text_step(image): | |
| extracted_text = extract_text(image) | |
| return extracted_text, extracted_text # Return extracted text and store it in state | |
| # Second step: Search and highlight keywords in the extracted text | |
| def highlight_keywords_step(extracted_text, keywords): | |
| highlighted_text = highlight_keywords(extracted_text, keywords) | |
| return highlighted_text | |
| # Gradio UI | |
| with gr.Blocks() as demo: | |
| # Step 1: Image Upload and Text Extraction | |
| 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") | |
| # Step 2: Keyword Input and Highlighting | |
| with gr.Row(): | |
| keyword_input = gr.Textbox(label="Enter keywords (comma-separated)", placeholder="Enter keywords after text extraction") | |
| search_button = gr.Button("Highlight Keywords") | |
| highlighted_text_output = gr.HTML(label="Highlighted Text with Matches") | |
| # Define interactions | |
| extract_button.click( | |
| fn=extract_text_step, # Call text extraction function | |
| inputs=image_input, | |
| outputs=[extracted_text_output, extracted_text_output], # Display text and store in state | |
| ) | |
| search_button.click( | |
| fn=highlight_keywords_step, # Call keyword highlighting function | |
| inputs=[extracted_text_output, keyword_input], # Use extracted text and keywords | |
| outputs=highlighted_text_output, # Display highlighted text | |
| ) | |
| # Launch the app | |
| demo.launch() | |