AtlasDgn/COVID_Example

1import streamlit as st2import numpy as np3from PIL import Image4import tensorflow as tf5from tensorflow.keras.models import load_model6 7# --- CONFIGURATION ---8# Define the three possible outcomes9CLASS_NAMES = {0: 'Normal', 1: 'Viral Pneumonia', 2: 'Covid'}10# The image size the model expects (224x224 pixels)11IMAGE_SIZE = 224 12 13# --- MODEL LOADING ---14# Use cache so the model only loads once15@st.cache_resource16def load_tuned_model():17    # Load the Keras model file. 18    # We include custom_objects to correctly load the VGG16 base.19    return tf.keras.models.load_model(20        "tuned_ai_model_best_lat.keras",21        custom_objects={'VGG16': tf.keras.applications.VGG16}22    )23 24# --- PREDICTION LOGIC ---25def run_prediction(image_file, model):26    """Processes the image and gets the diagnosis from the model."""27    try:28        # 1. Load and prepare the image29        image = Image.open(image_file).convert("RGB")30        img_array = np.array(image.resize((IMAGE_SIZE, IMAGE_SIZE)))31        32        # Add a dimension for the batch (1, 224, 224, 3)33        img_array = np.expand_dims(img_array, axis=0) 34        35        # Normalize pixel values (0 to 1)36        img_array = img_array / 255.037        38        # 2. Make prediction39        # The result is an array of probabilities for all three classes40        prediction_probabilities = model.predict(img_array).flatten()41        42        # 3. Find the most likely class43        class_index = np.argmax(prediction_probabilities)44        predicted_name = CLASS_NAMES[class_index]45        predicted_prob = prediction_probabilities[class_index]46        47        return predicted_name, predicted_prob48        49    except Exception as e:50        st.error(f"An error occurred during prediction: {e}")51        # Return None if any error happens52        return None, None53 54# --- STREAMLIT INTERFACE ---55 56st.title("COVID Detection from Chest X-ray")57st.markdown("Upload a chest X-ray image for diagnosis (Normal, Viral Pneumonia, or COVID).")58 59# Attempt to load the model and stop if it fails60try:61    model = load_tuned_model()62except Exception as e:63    st.error("Model Loading Failed. Please check dependencies and model file.")64    st.stop()65 66# --- UPLOAD SECTION ---67uploaded_file = st.file_uploader("Choose an X-ray image...", type=["jpg", "jpeg", "png"])68 69if uploaded_file is not None:70    # Display the uploaded image71    image = Image.open(uploaded_file)72    st.image(image, caption="Uploaded X-ray Image", use_container_width=True)73    74    # Run prediction when the button is clicked75    if st.button("Predict Diagnosis", type="primary"):76        77        # Run the prediction logic78        predicted_name, predicted_prob = run_prediction(uploaded_file, model)79        80        if predicted_name:81            st.markdown("---")82            st.subheader("Predicted Diagnosis")83            84            # Display the result simply (no emojis)85            if predicted_name == 'Covid':86                st.error(f"Result: **{predicted_name}**")87            else:88                st.success(f"Result: **{predicted_name}**")89            90            # Use the expander to show probability on click91            with st.expander(f"View Confidence Score for {predicted_name}"):92                st.markdown(f"Confidence: **{predicted_prob*100:.2f}%**")