Update app.py

app.py

@@ -19,7 +19,7 @@ transform = transforms.Compose([
 
 # Streamlit UI
 st.title("Saliva Disease Detection App")
-st.subheader("Predict Streptococcosis vs NOT Streptococcosis from uploaded saliva images")
+st.subheader("Predict Streptococcal infection present vs Streptococcal infection not present from uploaded saliva images")
 
 # Initialize session state for managing the uploaded file
 if "uploaded_file" not in st.session_state:
@@ -46,13 +46,14 @@ if st.session_state["uploaded_file"] is not None:
         _, predicted_class = torch.max(outputs, 1)
 
     # Map predicted class to labels
-    class_names = ['Not_Streptococcosis', 'Streptococcosis']
+    #class_names = ['Not_Streptococcosis', 'Streptococcosis']
+    class_names = ['Streptococcal_Infection_Not_Present', 'Streptococcal_Infection_Present']
     predicted_label = class_names[predicted_class.item()]
     predicted_probability = probabilities[0][predicted_class.item()].item() * 100  # Convert to percentage
 
     # Display the result
     st.write("### Prediction Result:")
-    if predicted_label == "Streptococcosis":
+    if predicted_label == "Streptococcal_Infection_Present":
         st.error(f"The sample is predicted as **{predicted_label}** with **{predicted_probability:.2f}%** probability.")
     else:
         st.success(f"The sample is predicted as **{predicted_label}** with **{predicted_probability:.2f}%** probability.")