Spaces:
Sleeping
Sleeping
| 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 = "" | |
| 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() | |