Update app.py

app.py

@@ -1,21 +1,21 @@
 import streamlit as st
 import tensorflow as tf
 import torch
+import torch.nn as nn
 import numpy as np
 from PIL import Image
 from torchvision import transforms, models
 
 torch_model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)
-torch_model.fc = torch.nn.Sequential(
-    torch.nn.Dropout(0.5),
-    torch.nn.Linear(2048, 256),
-    torch.nn.ReLU(),
-    torch.nn.Linear(256, 1),
+torch_model.fc = nn.Sequential(
+    nn.Dropout(0.5),
+    nn.Linear(2048, 256),
+    nn.ReLU(),
+    nn.Linear(256, 1),
 )
 
 tf_model = tf.keras.models.load_model('./models/tensorflow_model.keras')
 torch_model.load_state_dict(torch.load('./models/pytorch_model.pth', weights_only=True, map_location='cpu'))
-torch_model.to('cpu')
 torch_model.eval()
 
 selected_model = st.selectbox(
@@ -35,7 +35,7 @@ if uploaded_file is not None:
     image = Image.open(uploaded_file).convert("RGB")
 
     # Display Uploaded Image
-    st.image(image, caption="Uploaded X-ray", use_column_width=True)
+    st.image(image, caption="Uploaded X-ray", use_container_width=True)
 
     image_size = 256  # Match model input size
     if selected_model == "PyTorch":
@@ -62,6 +62,5 @@ if uploaded_file is not None:
         prediction = tf_model.predict(image_array)[0][0]  # Extract single value
 
     # Display Prediction
-    result = "Pneumonia Detected 🛑" if prediction > 0.5 else "No Pneumonia ✅"
-    st.write(f"**Prediction:** {result} ({prediction:.4f})")
+    st.write(f"**According to the model, there is a {prediction:.2f}% chance of having pneumonia!**")