Update app.py

app.py

@@ -1,5 +1,5 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import pipeline, AutoProcessor, AutoModelForCausalLM
 from PIL import Image
 import fitz  # PyMuPDF for PDF processing
 import logging
@@ -16,8 +16,10 @@ def setup_logging():
 @st.cache_resource
 def load_models():
     logging.info("Loading Hugging Face models...")
-    # Use most popular image-to-text model
-    image_to_text = pipeline("image-to-text", model="microsoft/trocr-large-printed")
+    
+    # Use a more reliable image-to-text model
+    image_to_text_processor = AutoProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+    image_to_text_model = AutoModelForCausalLM.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
     
     # Translation models
     translator_hi = pipeline("translation", model="Helsinki-NLP/opus-mt-en-hi")
@@ -26,16 +28,22 @@ def load_models():
     # Summarization model
     summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
     
-    return image_to_text, translator_hi, translator_ur, summarizer
+    return image_to_text_processor, image_to_text_model, translator_hi, translator_ur, summarizer
 
 # Function to extract text from images
-def extract_text_from_image(image):
+def extract_text_from_image(image, processor, model):
     logging.info("Extracting text from image...")
-    # Use TrOCR for more accurate text extraction
-    image_to_text = load_models()[0]
-    results = image_to_text(image)
-    # Combine all detected text
-    return " ".join([result['generated_text'] for result in results])
+    
+    # Prepare image for model
+    inputs = processor(images=image, return_tensors="pt")
+    
+    # Generate text
+    outputs = model.generate(**inputs)
+    
+    # Decode the generated text
+    preds = processor.decode(outputs[0], skip_special_tokens=True)
+    
+    return preds
 
 # Function to extract text from PDFs
 def extract_text_from_pdf(pdf_file):
@@ -61,7 +69,7 @@ def main():
     st.write("Upload a file (Image, PDF, or Text) to analyze and summarize the lab report in English, Hindi, and Urdu.")
     
     # Load all models
-    image_to_text, translator_hi, translator_ur, summarizer = load_models()
+    image_to_text_processor, image_to_text_model, translator_hi, translator_ur, summarizer = load_models()
     
     file = st.file_uploader("Upload a file (Image, PDF, or Text):", type=["jpg", "png", "jpeg", "pdf", "txt"])
     
@@ -70,7 +78,7 @@ def main():
         try:
             if file.type in ["image/jpeg", "image/png", "image/jpg"]:
                 image = Image.open(file)
-                text = extract_text_from_image(image)
+                text = extract_text_from_image(image, image_to_text_processor, image_to_text_model)
             elif file.type == "application/pdf":
                 text = extract_text_from_pdf(file)
             elif file.type == "text/plain":