Update app.py

app.py

@@ -1,111 +1,74 @@
 import os
 import streamlit as st
+from transformers import pipeline, AutoImageProcessor, AutoModelForImageClassification
 from PIL import Image
-import torch
-from torchvision import transforms, models
-import numpy as np
-from groq import Groq
 
-# Set up environment variables
-os.environ["GROQ_API_KEY"] = "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"
-
-# Initialize Groq client
-client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
-
-# Load Pretrained Models
-@st.cache_resource
-# Load Pretrained Model for Organ Recognition
-@st.cache_resource
-def load_organ_model():
-    model = models.resnet18(pretrained=True)  # Load pretrained ResNet18
-    num_features = model.fc.in_features       # Get the number of input features to the final layer
-    model.fc = torch.nn.Linear(num_features, 4)  # Modify the final layer for 4 classes
-    model.eval()  # Set the model to evaluation mode
-    return model
-
-
-
-# Image Preprocessing
-def preprocess_image(image):
-    transform = transforms.Compose([
-        transforms.Resize((224, 224)),
-        transforms.ToTensor(),
-        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
-    ])
-    return transform(image).unsqueeze(0)
-
-# Groq API for AI Insights
-def get_ai_insights(text_prompt):
+def load_pipeline():
+    """Load the Hugging Face pipeline for image classification."""
     try:
-        response = client.chat.completions.create(
-            messages=[{"role": "user", "content": text_prompt}],
-            model="llama-3.3-70b-versatile"
-        )
-        return response.choices[0].message.content
+        return pipeline("image-classification", model="dima806/pneumonia_chest_xray_image_detection")
     except Exception as e:
-        return f"Error: {e}"
-        
-# Organ Recognition Prediction
-def predict_organ(image):
-    with torch.no_grad():
-        input_tensor = preprocess_image(image)
-        output = organ_model(input_tensor)
-
-        # Check the output dimensions
-        st.write(f"Model output shape: {output.shape}")
-
-        # Ensure the output matches the number of classes
-        classes = ["Lungs", "Heart", "Spine", "Other"]
-        if output.size(1) != len(classes):
-            raise ValueError(
-                f"Model output size ({output.size(1)}) does not match the number of classes ({len(classes)})."
-            )
-
-        # Get the prediction
-        prediction_index = output.argmax().item()
-        prediction = classes[prediction_index]
-        return prediction
+        st.error(f"Error loading pipeline: {e}")
+        return None
 
-
-# Predict Normal/Abnormal
-def predict_normal_abnormal(image):
-    with torch.no_grad():
-        output = chexnet_model(preprocess_image(image))
-        classes = ["Normal", "Abnormal"]
-        prediction = classes[output.argmax().item()]
-    return prediction
+def classify_image_with_pipeline(pipe, image):
+    """Classify an image using the pipeline."""
+    try:
+        results = pipe(image)
+        return results
+    except Exception as e:
+        st.error(f"Error classifying image: {e}")
+        return None
 
 # Streamlit App
-st.title("Medical X-ray Analysis App")
-st.sidebar.title("Navigation")
-task = st.sidebar.radio("Select a task", ["Upload X-ray", "AI Insights"])
-
-if task == "Upload X-ray":
-    uploaded_file = st.file_uploader("Upload an X-ray image", type=["jpg", "png", "jpeg"])
-
-    if uploaded_file:
-        image = Image.open(uploaded_file)
-        st.image(image, caption="Uploaded X-ray", use_column_width=True)
-
-        # Predict Organ
-        st.subheader("Step 1: Identify the Organ")
-        organ = predict_organ(image)
-        st.write(f"Predicted Organ: **{organ}**")
-
-        # Predict Normal/Abnormal
-        st.subheader("Step 2: Analyze the X-ray")
-        classification = predict_normal_abnormal(image)
-        st.write(f"X-ray Status: **{classification}**")
-
-        if classification == "Abnormal":
-            st.subheader("Step 3: AI-Based Insights")
-            ai_prompt = f"Explain why this X-ray of the {organ} is abnormal."
-            insights = get_ai_insights(ai_prompt)
-            st.write(insights)
+st.title("Pneumonia Chest X-ray Image Detection")
+st.markdown(
+    """
+    This app detects signs of pneumonia in chest X-ray images using a pre-trained Hugging Face model.
+    """
+)
+
+# File uploader
+uploaded_file = st.file_uploader("Upload a chest X-ray image", type=["jpg", "jpeg", "png"])
+
+if uploaded_file:
+    image = Image.open(uploaded_file)
+    st.image(image, caption="Uploaded Chest X-ray", use_column_width=True)
+
+    # Load the model pipeline
+    pipe = load_pipeline()
+
+    if pipe:
+        st.write("Classifying the image...")
+        results = classify_image_with_pipeline(pipe, image)
+
+        if results:
+            st.write("### Classification Results:")
+            for result in results:
+                st.write(f"**Label:** {result['label']} | **Score:** {result['score']:.4f}")
+
+# Optional: Add Groq API integration if applicable
+if os.getenv("GROQ_API_KEY"):
+    from groq import Groq
+
+    client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
+
+    st.sidebar.markdown("### Groq API Integration")
+    question = st.sidebar.text_input("Ask a question about pneumonia or X-ray diagnosis:")
+
+    if question:
+        try:
+            chat_completion = client.chat.completions.create(
+                messages=[
+                    {
+                        "role": "user",
+                        "content": question,
+                    }
+                ],
+                model="llama-3.3-70b-versatile",
+            )
 
-elif task == "AI Insights":
-    st.subheader("Ask AI")
-    user_input = st.text_area("Enter your query for AI insights")
-    if user_input:
-        response = get_ai_insights(user_input)
-        st.write(response)
+            st.sidebar.write("**Groq API Response:**")
+            st.sidebar.write(chat_completion.choices[0].message.content)
+        except Exception as e:
+            st.sidebar.error(f"Error using Groq API: {e}")