Update app.py

app.py

@@ -1,70 +1,47 @@
 import streamlit as st
+import pickle
 from PIL import Image
-import torch
-from torchvision import models, transforms
 
-# Load the pre-trained model
-model = models.densenet121(pretrained=True)
-model.eval()
+# Load the pretrained model from the pickle file
+model_filename = 'model.pkl'
+with open(model_filename, 'rb') as file:
+    model = pickle.load(file)
 
-# Define the image transformations
-transform = transforms.Compose([
-    transforms.Resize(256),
-    transforms.CenterCrop(224),
-    transforms.ToTensor(),
-    transforms.Normalize(
-        mean=[0.485, 0.456, 0.406],
-        std=[0.229, 0.224, 0.225]
-    ),
-])
+# Function to make predictions
+def predict_pneumonia(image):
+    # Preprocess the image (you may need to resize or normalize it)
+    # preprocess_image(image)
 
-# Define the class labels
-class_labels = ['Normal', 'Pneumonia']
+    # Make predictions using the loaded model
+    prediction = model.predict(image)
 
-# Create a function to preprocess the image
-def preprocess_image(image):
-    # Convert the image to RGB
-    image = image.convert('RGB')
+    return prediction
 
-    # Resize the image to match the model's input shape
-    image = image.resize((224, 224))
-
-    # Convert the image to a tensor
-    image_tensor = transform(image)
-
-    # Add a batch dimension
-    image_tensor = image_tensor.unsqueeze(0)
+# Streamlit app
+def main():
+    # Set app title and layout
+    st.title("Pneumonia Detection")
+    st.markdown("---")
 
-    return image_tensor
+    # Add an image uploader
+    st.header("Upload Chest X-ray Image")
+    uploaded_file = st.file_uploader("Choose an image", type=["jpg", "jpeg", "png"])
 
-# Create a function to make predictions
-def predict(image):
-    # Preprocess the image
-    preprocessed_image = preprocess_image(image)
+    if uploaded_file is not None:
+        # Display the uploaded image
+        image = Image.open(uploaded_file)
+        st.image(image, caption="Uploaded Image", use_column_width=True)
 
-    # Make the prediction
-    with torch.no_grad():
-        output = model(preprocessed_image)
-        _, predicted_idx = torch.max(output, 1)
-        predicted_label = class_labels[predicted_idx.item()]
-    
-    return predicted_label
+        # Make prediction when the user clicks the 'Predict' button
+        if st.button("Predict"):
+            # Perform prediction
+            prediction = predict_pneumonia(image)
 
-# Create the Streamlit app
-def main():
-    st.title("Pneumonia Detection")
-    st.write("Upload an image and the app will predict if it has pneumonia or not.")
-    
-    # Upload and display the image
-    uploaded_image = st.file_uploader("Choose an image", type=["jpg", "jpeg", "png"])
-    
-    if uploaded_image is not None:
-        image = Image.open(uploaded_image)
-        st.image(image, caption="Uploaded Image", use_column_width=True)
-        
-        # Make a prediction
-        predicted_label = predict(image)
-        st.write("Prediction:", predicted_label)
+            # Display the prediction
+            if prediction == 1:
+                st.error("Prediction: Pneumonia detected")
+            else:
+                st.success("Prediction: No pneumonia detected")
 
 # Run the app
 if __name__ == '__main__':