Update app.py

app.py

@@ -1,84 +1,100 @@
 import os
 import streamlit as st
-import tensorflow as tf
-from tensorflow.keras.applications import ResNet50
-from tensorflow.keras.applications.resnet50 import preprocess_input, decode_predictions
-from tensorflow.keras.models import Model
-from tensorflow.keras.layers import Dense, GlobalAveragePooling2D
-from tensorflow.keras.preprocessing.image import img_to_array, load_img
-import numpy as np
-import requests
 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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"
 
-# Set Groq API key in environment variable
-os.environ['GROQ_API_KEY'] = "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"
-GROQ_API_KEY = os.getenv('GROQ_API_KEY')
-
-# Load pre-trained ResNet50 for normal/abnormal classification
-base_model = ResNet50(weights='imagenet', include_top=False, input_shape=(224, 224, 3))
-x = base_model.output
-x = GlobalAveragePooling2D()(x)
-x = Dense(1024, activation='relu')(x)
-predictions = Dense(1, activation='sigmoid')(x)
-classification_model = Model(inputs=base_model.input, outputs=predictions)
-
-# Load pre-trained ResNet50 for organ recognition
-organ_model = ResNet50(weights='imagenet')
-
-def classify_image(image_path):
-    """Classify the image as normal or abnormal."""
-    image = load_img(image_path, target_size=(224, 224))
-    image_array = img_to_array(image)
-    image_array = preprocess_input(image_array)
-    image_array = np.expand_dims(image_array, axis=0)
-    prediction = classification_model.predict(image_array)
-    return 'Abnormal' if prediction[0][0] > 0.5 else 'Normal'
-
-def recognize_organ(image_path):
-    """Recognize the organ in the image."""
-    image = load_img(image_path, target_size=(224, 224))
-    image_array = img_to_array(image)
-    image_array = preprocess_input(image_array)
-    image_array = np.expand_dims(image_array, axis=0)
-    prediction = organ_model.predict(image_array)
-    decoded = decode_predictions(prediction, top=3)[0]
-    return decoded[0][1]  # Top predicted class
-
-def get_ai_insights(organ):
-    """Fetch AI-based insights about the organ using Groq API."""
-    url = "https://api.groq.com/v1/insights"
-    headers = {"Authorization": f"Bearer {GROQ_API_KEY}", "Content-Type": "application/json"}
-    data = {"query": f"Provide detailed insights about {organ} X-ray, its diseases, and treatments."}
-    response = requests.post(url, headers=headers, json=data)
-    if response.status_code == 200:
-        return response.json().get("insights", "No insights available.")
-    else:
-        return "Failed to fetch insights. Please try again later."
-
-def main():
-    st.title("Medical Image Classification App")
-    st.sidebar.title("Navigation")
+# Initialize Groq client
+client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
+
+# Load Pretrained Models
+@st.cache_resource
+def load_model():
+    # Pretrained EfficientNet for organ recognition
+    organ_model = models.efficientnet_b0(pretrained=True)
+    organ_model.eval()
     
-    uploaded_file = st.file_uploader("Upload an X-ray or MRI image", type=["jpg", "jpeg", "png"])
+    # Pretrained DenseNet (CheXNet) for normal/abnormal classification
+    chexnet_model = models.densenet121(pretrained=True)
+    chexnet_model.classifier = torch.nn.Linear(1024, 2)  # Normal, Abnormal
+    chexnet_model.eval()
+
+    return organ_model, chexnet_model
+
+organ_model, chexnet_model = load_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):
+    try:
+        response = client.chat.completions.create(
+            messages=[{"role": "user", "content": text_prompt}],
+            model="llama-3.3-70b-versatile"
+        )
+        return response.choices[0].message.content
+    except Exception as e:
+        return f"Error: {e}"
+
+# Predict Organ
+def predict_organ(image):
+    with torch.no_grad():
+        output = organ_model(preprocess_image(image))
+        classes = ["Lungs", "Heart", "Spine", "Other"]  # Example classes
+        prediction = classes[output.argmax().item()]
+    return prediction
+
+# 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
+
+# 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 Image", use_column_width=True)
-        
-        with open("temp_image.jpg", "wb") as f:
-            f.write(uploaded_file.getbuffer())
+        st.image(image, caption="Uploaded X-ray", use_column_width=True)
 
-        st.write("### Classification Result")
-        result = classify_image("temp_image.jpg")
-        st.write(f"The X-ray is classified as: **{result}**")
+        # Predict Organ
+        st.subheader("Step 1: Identify the Organ")
+        organ = predict_organ(image)
+        st.write(f"Predicted Organ: **{organ}**")
 
-        st.write("### Organ Recognition")
-        organ = recognize_organ("temp_image.jpg")
-        st.write(f"Recognized 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}**")
 
-        st.write("### AI-Based Insights")
-        insights = get_ai_insights(organ)
-        st.write(insights)
+        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)
 
-if __name__ == "__main__":
-    main()
+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)