Spaces:
Sleeping
Sleeping
File size: 3,896 Bytes
fc0bd58 |
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 109 110 111 |
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()
|