Update app.py

app.py

@@ -1,44 +1,45 @@
 import streamlit as st
 from PIL import Image
-from huggingface_hub import InferenceClient
+from transformers import pipeline
 import io
 
-def analyze_image_with_maira(image):
-    """Analyzes the image using the Maira-2 model via the Hugging Face Inference API.
-    """
-    try:
-        # Prepare image data - send as bytes directly to 'inputs'
-        image_bytes = io.BytesIO()
-        if image.mode == "RGBA":  # Handle RGBA images (if any)
-            image = image.convert("RGB")
-        image.save(image_bytes, format="JPEG")
-        image_bytes = image_bytes.getvalue()  # Get the bytes
-
-        client = InferenceClient()  # No token needed inside the Space
-        result = client.post(
-            json={
-                "inputs": {
-                    "image": image_bytes.decode(encoding="latin-1"),  #Needs to be decoded
-                    "question": "Analyze this chest X-ray image and provide detailed findings. Include any abnormalities, their locations, and potential diagnoses. Be as specific as possible.",
-                            }
-            },
-            model="microsoft/maira-2",  # Specify the model
-            task="visual-question-answering"
-        )
-        return result
+# --- Configuration ---
+# Specify the model
+MODEL_NAME = "microsoft/maira-2"
+
+# --- Model Loading (using pipeline) ---
 
+@st.cache_resource  # Cache the pipeline for performance
+def load_pipeline():
+    """Loads the VQA pipeline."""
+    try:
+        # Explicitly set device if CUDA is available, otherwise use CPU
+        device = 0 if any(tf.config.list_physical_devices('GPU')) else -1
+        vqa_pipeline = pipeline("visual-question-answering", model=MODEL_NAME, device=device) # Add device
+        return vqa_pipeline
     except Exception as e:
-        st.error(f"An error occurred: {e}")
+        st.error(f"Error loading pipeline: {e}")
         return None
 
+# --- Image Preprocessing (Keep as bytes) ---
+def prepare_image(image):
+    """Prepares the PIL Image object for the pipeline (handles RGBA)."""
+    image_bytes = io.BytesIO()
+    if image.mode == "RGBA":
+        image = image.convert("RGB")
+    image.save(image_bytes, format="JPEG")
+    return image_bytes.getvalue() # Return bytes directly
 
 # --- Streamlit App ---
 
 def main():
-    st.title("Chest X-ray Analysis with Maira-2 (Hugging Face Spaces)")
-    st.write(
-        "Upload a chest X-ray image. This app uses the Maira-2 model within this Hugging Face Space."
-    )
+    st.title("Chest X-ray Analysis with Maira-2 (Transformers Pipeline)")
+    st.write("Upload a chest X-ray image.  This app uses the Maira-2 model via the Transformers library.")
+
+    vqa_pipeline = load_pipeline()
+    if vqa_pipeline is None:
+      st.warning("Pipeline not loaded.  Predictions will not be available.")
+      return
 
     uploaded_file = st.file_uploader("Choose a chest X-ray image (JPG, PNG)", type=["jpg", "jpeg", "png"])
 
@@ -47,18 +48,26 @@ def main():
         st.image(image, caption="Uploaded Image", use_column_width=True)
 
         with st.spinner("Analyzing image with Maira-2..."):
-            analysis_results = analyze_image_with_maira(image)
-
-        if analysis_results:
-            # --- Results Display (VQA format) ---
-            if isinstance(analysis_results, dict) and 'answer' in analysis_results:
-                st.subheader("Findings:")
-                st.write(analysis_results['answer'])
-            else:
-                st.warning("Unexpected API response format.")
-                st.write("Raw API response:", analysis_results)
-        else:
-            st.error("Failed to get analysis results.")
+            image_data = prepare_image(image)
+            try:
+                results = vqa_pipeline(
+                    image=image_data,  # Pass the image bytes
+                    question="Analyze this chest X-ray image and provide detailed findings. Include any abnormalities, their locations, and potential diagnoses. Be as specific as possible.",
+                )
+
+                if results:  # Handle results (list of dicts)
+                  if isinstance(results, list) and len(results) > 0:
+                      best_answer = max(results, key=lambda x: x.get('score', 0))
+                      if 'answer' in best_answer:
+                          st.subheader("Findings:")
+                          st.write(best_answer['answer'])
+                      else:
+                          st.warning("Could not find 'answer' in results.")
+                  else:
+                      st.warning("Unexpected result format.")
+
+            except Exception as e:
+                st.error(f"An error occurred during analysis: {e}")
 
     else:
         st.write("Please upload an image.")
@@ -66,5 +75,6 @@ def main():
     st.write("---")
     st.write("Disclaimer: For informational purposes only.  Not medical advice.")
 
+
 if __name__ == "__main__":
     main()