Update app.py

app.py

@@ -3,84 +3,90 @@ import numpy as np
 from PIL import Image
 import tensorflow as tf
 from tensorflow.keras.models import load_model
-import os
-import time
-from io import BytesIO
 
-# Define the correct class names for output
+# --- CONFIGURATION ---
+# Define the three possible outcomes
 CLASS_NAMES = {0: 'Normal', 1: 'Viral Pneumonia', 2: 'Covid'}
-IMG_SIZE = 224 
+# The image size the model expects (224x224 pixels)
+IMAGE_SIZE = 224 
 
-# Function to load the model (cached for efficiency)
+# --- MODEL LOADING ---
+# Use cache so the model only loads once
 @st.cache_resource
 def load_tuned_model():
-    # Use custom_objects for maximum robustness when loading models with 
-    # pre-trained components like VGG16 in specific environments.
+    # Load the Keras model file. 
+    # We include custom_objects to correctly load the VGG16 base.
     return tf.keras.models.load_model(
         "tuned_ai_model_best_lat.keras",
         custom_objects={'VGG16': tf.keras.applications.VGG16}
     )
 
-# Function to run the prediction and show progress
-def run_prediction(image_file, model, img_size):
-    # Progress visualization (for user experience)
-    progress_bar = st.progress(0)
-    status_text = st.empty()
-    for i in range(100):
-        progress_bar.progress(i + 1)
-        status_text.text(f"Processing... {i+1}%")
-        time.sleep(0.01)
-
+# --- PREDICTION LOGIC ---
+def run_prediction(image_file, model):
+    """Processes the image and gets the diagnosis from the model."""
     try:
-        # Read and process the image using PIL (Image)
+        # 1. Load and prepare the image
         image = Image.open(image_file).convert("RGB")
-        img_array = np.array(image.resize((img_size, img_size)))
-        img_array = np.expand_dims(img_array, axis=0)
+        img_array = np.array(image.resize((IMAGE_SIZE, IMAGE_SIZE)))
+        
+        # Add a dimension for the batch (1, 224, 224, 3)
+        img_array = np.expand_dims(img_array, axis=0) 
+        
+        # Normalize pixel values (0 to 1)
         img_array = img_array / 255.0
-
-        # Make the prediction
-        prediction = model.predict(img_array).flatten()
-
-        # Find the predicted class index and name
-        class_predicted_idx = np.argmax(prediction)
-        predicted_class_name = CLASS_NAMES[class_predicted_idx]
-
-        status_text.success("Prediction complete!")
         
-        # Display results
-        st.subheader("Prediction Probabilities:")
-        for i, prob in enumerate(prediction):
-            class_name = CLASS_NAMES[i]
-            if i == class_predicted_idx:
-                st.markdown(f"**{class_name}: {prob*100:.2f}%** (Predicted)", unsafe_allow_html=True)
-            else:
-                st.markdown(f"{class_name}: {prob*100:.2f}%")
-
-        if class_predicted_idx == 2:
-            st.error(f"**Result: ❗ HIGH LIKELIHOOD OF COVID DETECTED ❗**")
-        else:
-            st.success(f"**Result: ✅ {predicted_class_name} Detected**")
+        # 2. Make prediction
+        # The result is an array of probabilities for all three classes
+        prediction_probabilities = model.predict(img_array).flatten()
+        
+        # 3. Find the most likely class
+        class_index = np.argmax(prediction_probabilities)
+        predicted_name = CLASS_NAMES[class_index]
+        predicted_prob = prediction_probabilities[class_index]
+        
+        return predicted_name, predicted_prob
         
     except Exception as e:
-        status_text.error(f"An error occurred during prediction: {e}")
-        st.exception(e) # Show full traceback for debugging
+        st.error(f"An error occurred during prediction: {e}")
+        # Return None if any error happens
+        return None, None
 
-# Load the model and ensure the app stops if loading fails
+# --- STREAMLIT INTERFACE ---
+
+st.title("COVID Detection from Chest X-ray")
+st.markdown("Upload a chest X-ray image for diagnosis (Normal, Viral Pneumonia, or COVID).")
+
+# Attempt to load the model and stop if it fails
 try:
     model = load_tuned_model()
 except Exception as e:
-    st.error("Model Loading Failed. Check TF/Keras versions and model file name (tuned_ai_model_best_lat.keras).")
-    st.exception(e)
+    st.error("Model Loading Failed. Please check dependencies and model file.")
     st.stop()
 
-st.title("COVID Detection from Chest X-ray")
-st.write("Upload a chest X-ray image to predict if it shows signs of Normal, Viral Pneumonia, or COVID.")
-
+# --- UPLOAD SECTION ---
 uploaded_file = st.file_uploader("Choose an X-ray image...", type=["jpg", "jpeg", "png"])
 
 if uploaded_file is not None:
+    # Display the uploaded image
     image = Image.open(uploaded_file)
     st.image(image, caption="Uploaded X-ray Image", use_container_width=True)
-
-    if st.button("Predict"):
-        run_prediction(uploaded_file, model, IMG_SIZE)
+    
+    # Run prediction when the button is clicked
+    if st.button("Predict Diagnosis", type="primary"):
+        
+        # Run the prediction logic
+        predicted_name, predicted_prob = run_prediction(uploaded_file, model)
+        
+        if predicted_name:
+            st.markdown("---")
+            st.subheader("Predicted Diagnosis")
+            
+            # Display the result simply (no emojis)
+            if predicted_name == 'Covid':
+                st.error(f"Result: **{predicted_name}**")
+            else:
+                st.success(f"Result: **{predicted_name}**")
+            
+            # Use the expander to show probability on click
+            with st.expander(f"View Confidence Score for {predicted_name}"):
+                st.markdown(f"Confidence: **{predicted_prob*100:.2f}%**")