Update app.py

app.py

@@ -1,68 +1,36 @@
-# -*- coding: utf-8 -*-
-"""OCR Web Application Prototype.ipynb
-Automatically generated by Colab.
-Original file is located at
-    https://colab.research.google.com/drive/1vzsQ17-W1Vy6yJ60XUwFy0QRkOR_SIg7
-"""
-
-from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
-from qwen_vl_utils import process_vision_info
+from transformers import AutoProcessor
 import torch
 import gradio as gr
 from PIL import Image
+from byaldi import RAGMultiModalModel
+from qwen_vl_utils import process_vision_info
 
+# Load ColPali model
+RAG = RAGMultiModalModel.from_pretrained("vidore/colpali")
 
 processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct")
 
-# Initialize the model with float16 precision and handle fallback to CPU
-# Simplified model loading function for CPU
 def load_model():
-    return Qwen2VLForConditionalGeneration.from_pretrained(
-        "Qwen/Qwen2-VL-2B-Instruct",
-        torch_dtype=torch.float32,  # Use float32 for CPU
-        low_cpu_mem_usage=True
-    )
+    return RAG.model
 
-# Load the model
 vlm = load_model()
 
-# OCR function to extract text from an image
-def ocr_image(image, query="Extract text from the image", keyword=""):
-    messages = [
-        {
-            "role": "user",
-            "content": [
-                {
-                    "type": "image",
-                    "image": image,
-                },
-                {"type": "text", "text": query},
-            ],
-        }
-    ]
+def ocr_image(image, keyword=""):
+    # Convert PIL Image to file-like object
+    import io
+    img_byte_arr = io.BytesIO()
+    image.save(img_byte_arr, format='PNG')
+    img_byte_arr = img_byte_arr.getvalue()
 
-    # Prepare inputs for the model
-    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    image_inputs, video_inputs = process_vision_info(messages)
-    inputs = processor(
-        text=[text],
-        images=image_inputs,
-        videos=video_inputs,
-        padding=True,
-        return_tensors="pt",
-    )
-    inputs = inputs.to("cpu")
+    # Index the image
+    RAG.index(input_data=img_byte_arr, index_name="temp_index", overwrite=True)
 
-    # Generate the output text using the model
-    generated_ids = vlm.generate(**inputs, max_new_tokens=512)
-    generated_ids_trimmed = [
-        out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
-    ]
-    
-    output_text = processor.batch_decode(
-        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
-    )[0]
+    # Retrieve text from the image
+    results = RAG.search("Extract all text from this image", k=1)
     
+    # Extract text from results
+    output_text = results[0].get('text', '')
+
     if keyword:
         keyword_lower = keyword.lower()
         if keyword_lower in output_text.lower():
@@ -73,14 +41,12 @@ def ocr_image(image, query="Extract text from the image", keyword=""):
     else:
         return output_text
 
-# Gradio interface
 def process_image(image, keyword=""):
     max_size = 1024
     if max(image.size) > max_size:
         image.thumbnail((max_size, max_size))
     return ocr_image(image, keyword=keyword)
 
-# Update the Gradio interface:
 interface = gr.Interface(
     fn=process_image,
     inputs=[
@@ -91,5 +57,4 @@ interface = gr.Interface(
     title="Hindi & English OCR with Keyword Search",
 )
 
-# Launch Gradio interface in Colab
 interface.launch()