Spaces:
Build error
Build error
File size: 3,619 Bytes
2f3144c 3ef82d2 2f3144c 7508b05 2f3144c 3ef82d2 7508b05 3ef82d2 2f3144c cacc570 7508b05 cacc570 2f3144c 7508b05 2f3144c 7508b05 2f3144c 7508b05 2f3144c 7508b05 2f3144c 7508b05 2be66f7 cacc570 2be66f7 7508b05 cacc570 7508b05 |
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 |
import gradio as gr
from byaldi import RAGMultiModalModel
from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info
import torch
from PIL import Image
import os
import traceback
import spaces
# Load the models
rag_model = RAGMultiModalModel.from_pretrained("vidore/colpali")
qwen_model = Qwen2VLForConditionalGeneration.from_pretrained(
"Qwen/Qwen2-VL-7B-Instruct", trust_remote_code=True, torch_dtype=torch.bfloat16
)
processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct", trust_remote_code=True)
# Global variable to store extracted text
extracted_text = ""
@spaces.GPU(duration=120)
def ocr_and_extract(image, text_query):
global extracted_text
try:
temp_image_path = "temp_image.jpg"
image.save(temp_image_path)
rag_model.index(input_path=temp_image_path, index_name="image_index", store_collection_with_index=False, overwrite=True)
results = rag_model.search(text_query, k=1)
image_data = Image.open(temp_image_path)
messages = [
{"role": "user", "content": [{"type": "image", "image": image_data}, {"type": "text", "text": text_query}]}
]
text_input = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
image_inputs, _ = process_vision_info(messages)
inputs = processor(text=[text_input], images=image_inputs, padding=True, return_tensors="pt")
qwen_model.to("cuda")
inputs = {k: v.to("cuda") for k, v in inputs.items()}
generated_ids = qwen_model.generate(**inputs, max_new_tokens=50)
output_text = processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
extracted_text = output_text[0]
os.remove(temp_image_path)
return extracted_text
except Exception as e:
traceback.print_exc()
return f"Error: {str(e)}"
def keyword_search(keywords):
if extracted_text:
found_keywords = [word for word in keywords.split() if word in extracted_text]
if found_keywords:
return f"Keywords found: {', '.join(found_keywords)}"
else:
return "No matching keywords found."
else:
return "No text extracted yet. Please upload an image."
# Interface Layout
extract_text_button = gr.Button("Extract Text")
extracted_text_box = gr.Textbox(label="Extracted Text", placeholder="Text will appear here...", interactive=False)
keyword_search_box = gr.Textbox(label="Enter keywords to search", placeholder="Type keywords here...")
search_results = gr.Textbox(label="Search Results", interactive=False)
# Re-order the components: Extract Text button goes above Extracted Text box
iface = gr.Interface(
fn=ocr_and_extract,
inputs=[gr.Image(type="pil"), gr.Textbox(label="Enter your query (optional)")],
outputs=[extracted_text_box],
title="Image OCR with Byaldi + Qwen2-VL",
description="Upload an image (JPEG/PNG) containing Hindi and English text for OCR."
)
# Layout for keyword search
search_interface = gr.Interface(
fn=keyword_search,
inputs=[keyword_search_box],
outputs=[search_results],
title="Keyword Search within Extracted Text",
description="Enter keywords to search within the extracted text."
)
# Combining both interfaces with keyword search on the same page
combined_interface = gr.Blocks()
with combined_interface:
extract_text_button.render()
extracted_text_box.render()
keyword_search_box.render()
search_results.render()
combined_interface.launch()
|