Update app.py

app.py

@@ -28,7 +28,7 @@ def load_model():
 vlm = load_model()
 
 # OCR function to extract text from an image
-def ocr_image(image, query="Extract text from the image"):
+def ocr_image(image, query="Extract text from the image", keyword=""):
     messages = [
         {
             "role": "user",
@@ -59,22 +59,34 @@ def ocr_image(image, query="Extract text from the image"):
     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
-    )
-    return output_text[0]
+    )[0]
+    
+    if keyword:
+        keyword_lower = keyword.lower()
+        if keyword_lower in output_text.lower():
+            highlighted_text = output_text.replace(keyword, f"**{keyword}**")
+            return f"Keyword '{keyword}' found in the text:\n\n{highlighted_text}"
+        else:
+            return f"Keyword '{keyword}' not found in the text:\n\n{output_text}"
+    else:
+        return output_text
 
 # Gradio interface
-def process_image(image):
-    return ocr_image(image)
+def process_image(image, keyword=""):
+    return ocr_image(image, keyword=keyword)
 
-# Create Gradio interface for uploading an image
+# Update the Gradio interface:
 interface = gr.Interface(
     fn=process_image,
-    inputs=gr.Image(type="pil"),
+    inputs=[
+        gr.Image(type="pil"),
+        gr.Textbox(label="Enter keyword to search (optional)")
+    ],
     outputs="text",
-    title="Hindi & English OCR",
-    description="Upload an image containing text in Hindi or English to extract the text using OCR."
+    title="Hindi & English OCR with Keyword Search",
 )
 
 # Launch Gradio interface in Colab