Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from transformers import pipeline, CLIPProcessor, CLIPModel
+from transformers import pipeline, CLIPProcessor, CLIPModel, VisionEncoderDecoderModel, TrOCRProcessor
 from PIL import Image
 import torch
 
@@ -11,6 +11,10 @@ captioner = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captionin
 clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
 clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
 
+# 3. New Strategy: Optical Character Recognition (OCR)
+ocr_processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-printed")
+ocr_model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-printed")
+
 # --- Functions ---
 def multi_agent_inference(image, search_text=""):
     outputs = []
@@ -44,21 +48,29 @@ def multi_agent_inference(image, search_text=""):
     except Exception as e:
         outputs.append(("Image-in-Image Search", f"[failed] {e}"))
 
+    # Strategy 4: Optical Character Recognition (OCR) 
+    try:
+        pixel_values = ocr_processor(images=image, return_tensors="pt").pixel_values
+        generated_ids = ocr_model.generate(pixel_values)
+        extracted_text = ocr_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        outputs.append(("OCR (Extracted Text)", extracted_text))
+    except Exception as e:
+        outputs.append(("OCR (Extracted Text)", f"[failed] {e}"))
+
     return outputs
 
 # --- Gradio UI ---
 with gr.Blocks() as demo:
     gr.Markdown("## 🤗 Multi-Strategy Hugging Face AI MVP")
-    gr.Markdown("Upload an image → get captions, search for text inside it, compare images. "
-                "All strategies use Hugging Face free OSS models.")
-
+    gr.Markdown("Upload an image → get captions, search for text inside it, "
+                "compare images, and **extract text**. All strategies use "
+                "Hugging Face free OSS models.")
     with gr.Row():
         img_input = gr.Image(type="pil", label="Upload Image")
-        text_input = gr.Textbox(label="Search Text (optional)", placeholder="e.g. cat, car, chip, war...")
+        text_input = gr.Textbox(label="Search Text (optional)", 
+                                placeholder="e.g. cat, car, chip, war...")
     btn = gr.Button("Run Multi-Agent Analysis")
-
     out = gr.Dataframe(headers=["Strategy", "Result"], label="🔍 AI Outputs")
-
     btn.click(fn=multi_agent_inference, inputs=[img_input, text_input], outputs=out)
 
-demo.launch()
+demo.launch()