Pneumonia-Detection / src /streamlit_app.py
DD009's picture
Update src/streamlit_app.py
d39c31e verified
import streamlit as st
import numpy as np
import tensorflow as tf
import pydicom
import cv2
# --------------------------------------------------
# App Config
# --------------------------------------------------
st.set_page_config(page_title="Pneumonia Detection", layout="centered")
st.title("🫁 Pneumonia Detection from Chest X-ray")
st.write("Upload a chest X-ray (DICOM) to detect pneumonia.")
# --------------------------------------------------
# Load Serialized Model (Inference Only)
# --------------------------------------------------
@st.cache_resource
def load_model():
return tf.keras.models.load_model(
"pneumonia_final_model.keras",
compile=False # avoids optimizer warnings
)
model = load_model()
# --------------------------------------------------
# Preprocess DICOM
# --------------------------------------------------
def preprocess_dicom(uploaded_file):
dicom = pydicom.dcmread(uploaded_file)
img = dicom.pixel_array.astype(np.float32)
img = (img - img.min()) / (img.max() - img.min() + 1e-6)
img = cv2.resize(img, (224, 224))
# grayscale β†’ RGB
img = np.stack([img] * 3, axis=-1)
img = np.expand_dims(img, axis=0)
return img
# --------------------------------------------------
# File Upload
# --------------------------------------------------
uploaded_file = st.file_uploader(
"Upload Chest X-ray (.dcm file)",
type=["dcm"]
)
# --------------------------------------------------
# Prediction
# --------------------------------------------------
if uploaded_file is not None:
try:
image = preprocess_dicom(uploaded_file)
with st.spinner("Analyzing X-ray..."):
probs = model.predict(image)[0]
# ---- 3-class β†’ binary mapping ----
pneumonia_prob = float(probs[2])
non_pneumonia_prob = float(probs[0] + probs[1])
if pneumonia_prob >= 0.5:
prediction = "Pneumonia"
confidence = pneumonia_prob
st.error(f"⚠️ {prediction} detected")
else:
prediction = "Non-Pneumonia"
confidence = non_pneumonia_prob
st.success(f"βœ… {prediction}")
st.subheader("Prediction Result")
st.write(f"**Prediction:** {prediction}")
st.write(f"**Confidence:** {confidence:.2%}")
# Optional: show raw probabilities (for transparency)
with st.expander("View raw class probabilities"):
st.write({
"Normal": float(probs[0]),
"No Lung Opacity": float(probs[1]),
"Lung Opacity (Pneumonia)": float(probs[2])
})
except Exception as e:
st.error(f"Error processing file: {e}")