Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,66 +1,228 @@
 import io
 import os
+import json
+from datetime import datetime
+
 import numpy as np
+import pandas as pd
 import streamlit as st
 import tensorflow as tf
 from tensorflow import keras
 import pydicom
+import matplotlib.pyplot as plt
+
+from fpdf import FPDF
 
-# ----------------------------------------------------
-# App Configuration
-# ----------------------------------------------------
+# -----------------------------
+# Page config
+# -----------------------------
 st.set_page_config(
     page_title="Pneumonia Detection (Chest X-ray) – Clinical Decision Support",
     layout="centered"
 )
 
-
 st.title("Pneumonia Detection (Chest X-ray) – Clinical Decision Support")
 st.caption(
-    "Upload one or more Chest X-ray DICOM files (.dcm). "
-    "Adjust the decision threshold and submit to obtain a probability-based binary prediction. "
-    "This system is intended for clinical decision support and does not replace professional medical judgment."
+    "Upload one or more Chest X-ray DICOM files (.dcm). Adjust the decision threshold and click Submit. "
+    "This tool is for decision support only and does not replace clinical judgment."
 )
 
-# ----------------------------------------------------
-# Load Model
-# ----------------------------------------------------
-
-
-MODEL_PATH = os.path.join(os.path.dirname(__file__), "..", "model.keras")
+# -----------------------------
+# Paths / Model Loading
+# -----------------------------
+REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+MODEL_PATH = os.path.join(REPO_ROOT, "model.keras")
 
+# Optional: store a version tag manually in a json file in repo root if you want
+VERSION_PATH = os.path.join(REPO_ROOT, "model_version.json")
 
 @st.cache_resource
 def load_model():
+    if not os.path.exists(MODEL_PATH):
+        raise FileNotFoundError(f"model.keras not found at: {MODEL_PATH}")
+
     try:
-        return keras.models.load_model(MODEL_PATH)
+        m = keras.models.load_model(MODEL_PATH)
     except Exception:
         keras.config.enable_unsafe_deserialization()
-        return keras.models.load_model(MODEL_PATH, safe_mode=False)
+        m = keras.models.load_model(MODEL_PATH, safe_mode=False)
+    return m
 
 model = load_model()
 
-input_shape = model.input_shape
+# read model input details
+input_shape = model.input_shape  # (None, H, W, C)
 img_size = int(input_shape[1]) if input_shape and input_shape[1] else 256
-expected_channels = int(input_shape[-1]) if input_shape and input_shape[-1] else 3
+exp_ch = int(input_shape[-1]) if input_shape and input_shape[-1] else 1
+
+# -----------------------------
+# Utilities
+# -----------------------------
+def get_model_version():
+    if os.path.exists(VERSION_PATH):
+        try:
+            with open(VERSION_PATH, "r") as f:
+                return json.load(f).get("version", "unknown")
+        except Exception:
+            return "unknown"
+    return "v1"
+
+MODEL_VERSION = get_model_version()
+
+def read_dicom(uploaded_file) -> np.ndarray:
+    data = uploaded_file.read()
+    dcm = pydicom.dcmread(io.BytesIO(data))
+    img = dcm.pixel_array.astype(np.float32)
+
+    # Normalize to 0..1
+    img_min = float(np.min(img))
+    img_max = float(np.max(img))
+    img = (img - img_min) / (img_max - img_min + 1e-8)
+
+    return img
+
+def preprocess(img_2d: np.ndarray) -> np.ndarray:
+    # (H,W) -> (1,H,W,C) float32 0..1
+    x = tf.convert_to_tensor(img_2d[..., np.newaxis], dtype=tf.float32)  # (H,W,1)
+    x = tf.image.resize(x, (img_size, img_size))
+    x = tf.clip_by_value(x, 0.0, 1.0)
+    x = x.numpy()
+
+    if exp_ch == 3 and x.shape[-1] == 1:
+        x = np.repeat(x, 3, axis=-1)
+    elif exp_ch == 1 and x.shape[-1] == 3:
+        x = x[..., :1]
+
+    x = np.expand_dims(x, axis=0)
+    return x.astype(np.float32)
+
+def predict_prob(x: np.ndarray) -> float:
+    pred = model.predict(x, verbose=0)
+    if isinstance(pred, (list, tuple)):
+        prob = float(np.ravel(pred[-1])[0])
+    else:
+        prob = float(np.ravel(pred)[0])
+    return max(0.0, min(1.0, prob))
+
+def confidence_bucket(prob: float) -> str:
+    # Clinical-friendly interpretation (you can adjust the bands)
+    if prob < 0.30:
+        return "Low likelihood (< 0.30)"
+    elif prob <= 0.60:
+        return "Borderline suspicion (0.30 – 0.60)"
+    else:
+        return "High likelihood (> 0.60)"
+
+# -----------------------------
+# Grad-CAM (ResNet-style) helper
+# -----------------------------
+def find_last_conv_layer(m: keras.Model) -> str:
+    # picks the last Conv2D layer name
+    for layer in reversed(m.layers):
+        if isinstance(layer, keras.layers.Conv2D):
+            return layer.name
+    # If model is nested and last conv is inside base model:
+    for layer in reversed(m.layers):
+        if isinstance(layer, keras.Model):
+            for sub in reversed(layer.layers):
+                if isinstance(sub, keras.layers.Conv2D):
+                    return sub.name
+    raise ValueError("Could not find a Conv2D layer for Grad-CAM.")
+
+@st.cache_resource
+def get_gradcam_model(m: keras.Model):
+    last_conv = find_last_conv_layer(m)
+    conv_layer = m.get_layer(last_conv)
+    grad_model = keras.Model([m.inputs], [conv_layer.output, m.output])
+    return grad_model, last_conv
+
+def make_gradcam_heatmap(x_input: np.ndarray) -> np.ndarray:
+    grad_model, _ = get_gradcam_model(model)
+
+    x_tensor = tf.convert_to_tensor(x_input, dtype=tf.float32)
+    with tf.GradientTape() as tape:
+        conv_out, preds = grad_model(x_tensor)
+
+        if isinstance(preds, (list, tuple)):
+            preds = preds[-1]
+
+        # binary prob is preds[:,0]
+        score = preds[:, 0]
 
-# ----------------------------------------------------
-# Threshold Slider (DEFAULT = 0.37 for ResNet)
-# ----------------------------------------------------
+    grads = tape.gradient(score, conv_out)
+    pooled = tf.reduce_mean(grads, axis=(0, 1, 2))
+    conv_out = conv_out[0]
+
+    heatmap = conv_out @ pooled[..., tf.newaxis]
+    heatmap = tf.squeeze(heatmap)
+
+    heatmap = tf.maximum(heatmap, 0)
+    denom = tf.reduce_max(heatmap) + 1e-8
+    heatmap = heatmap / denom
+    return heatmap.numpy()
+
+def overlay_heatmap_on_image(img_2d: np.ndarray, heatmap: np.ndarray):
+    # Resize heatmap to img_size
+    heat = tf.image.resize(heatmap[..., None], (img_size, img_size)).numpy().squeeze()
+
+    fig = plt.figure(figsize=(5, 5))
+    plt.imshow(img_2d, cmap="gray")
+    plt.imshow(heat, cmap="jet", alpha=0.35)
+    plt.axis("off")
+    plt.tight_layout()
+    return fig
+
+# -----------------------------
+# PDF generator
+# -----------------------------
+def build_pdf_report(df: pd.DataFrame, threshold: float) -> bytes:
+    pdf = FPDF()
+    pdf.add_page()
+    pdf.set_font("Arial", size=12)
+
+    pdf.multi_cell(0, 8, f"Pneumonia Detection Report")
+    pdf.ln(1)
+    pdf.set_font("Arial", size=10)
+    pdf.multi_cell(0, 6, f"Generated at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
+    pdf.multi_cell(0, 6, f"Model version: {MODEL_VERSION}")
+    pdf.multi_cell(0, 6, f"Decision threshold used: {threshold:.2f}")
+    pdf.ln(2)
+
+    # Table header
+    pdf.set_font("Arial", "B", 9)
+    headers = ["file_name", "probability", "prediction", "confidence_band"]
+    col_w = [70, 25, 35, 55]
+    for h, w in zip(headers, col_w):
+        pdf.cell(w, 7, h, border=1)
+    pdf.ln()
+
+    # Rows
+    pdf.set_font("Arial", size=9)
+    for _, r in df.iterrows():
+        pdf.cell(col_w[0], 7, str(r["file_name"])[:40], border=1)
+        pdf.cell(col_w[1], 7, f'{float(r["probability"]):.4f}', border=1)
+        pdf.cell(col_w[2], 7, str(r["prediction"])[:18], border=1)
+        pdf.cell(col_w[3], 7, str(r["confidence_band"])[:30], border=1)
+        pdf.ln()
+
+    return pdf.output(dest="S").encode("latin-1")
+
+# -----------------------------
+# UI
+# -----------------------------
 st.subheader("Model Parameters")
 
 threshold = st.slider(
     "Decision Threshold",
     min_value=0.01,
     max_value=0.99,
-    value=0.37,  
+    value=0.37,   # your ResNet best-thr default
     step=0.01,
-    help="If predicted probability ≥ threshold → Pneumonia. Otherwise → Not Pneumonia."
+    help="If predicted probability ≥ threshold → Pneumonia, else → Not Pneumonia."
 )
 
-# ----------------------------------------------------
-# File Upload
-# ----------------------------------------------------
+show_gradcam = st.checkbox("Show Grad-CAM heatmap (explainability)", value=True)
+
 st.subheader("Upload Chest X-ray DICOM Files")
 
 uploaded_files = st.file_uploader(
@@ -76,74 +238,113 @@ with col2:
     clear = st.button("Clear", use_container_width=True)
 
 if clear:
-    st.experimental_rerun()
+    st.rerun()
 
-# ----------------------------------------------------
-# Helper Functions
-# ----------------------------------------------------
-def read_dicom(file):
-    data = file.read()
-    dcm = pydicom.dcmread(io.BytesIO(data))
-    img = dcm.pixel_array.astype(np.float32)
+st.subheader("Prediction Results")
 
-    img = (img - img.min()) / (img.max() - img.min() + 1e-8)
-    return img
+if submit:
+    if not uploaded_files:
+        st.warning("Please upload at least one DICOM file before submitting.")
+    else:
+        rows = []
+        with st.spinner("Running inference..."):
+            for f in uploaded_files:
+                try:
+                    img = read_dicom(f)
+                    x = preprocess(img)
+                    prob = predict_prob(x)
+                    pred_label = "Pneumonia" if prob >= threshold else "Not Pneumonia"
+                    band = confidence_bucket(prob)
 
-def preprocess(img):
-    x = tf.convert_to_tensor(img[..., None], dtype=tf.float32)
-    x = tf.image.resize(x, (img_size, img_size))
-    x = tf.clip_by_value(x, 0.0, 1.0)
-    x = x.numpy()
+                    rows.append({
+                        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+                        "model_version": MODEL_VERSION,
+                        "file_name": f.name,
+                        "probability": prob,
+                        "prediction": pred_label,
+                        "confidence_band": band
+                    })
 
-    # If model expects 3 channels (ResNet)
-    if expected_channels == 3 and x.shape[-1] == 1:
-        x = np.repeat(x, 3, axis=-1)
+                except Exception as e:
+                    rows.append({
+                        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+                        "model_version": MODEL_VERSION,
+                        "file_name": f.name,
+                        "probability": np.nan,
+                        "prediction": "Error",
+                        "confidence_band": str(e)
+                    })
 
-    x = np.expand_dims(x, axis=0)
-    return x.astype(np.float32)
+        df = pd.DataFrame(rows)
 
-def get_probability(x):
-    prediction = model.predict(x, verbose=0)
+        # Sentence-style outputs
+        for _, r in df.iterrows():
+            if r["prediction"] == "Error":
+                st.error(
+                    f"For the uploaded file '{r['file_name']}', the system could not generate a prediction. "
+                    f"Reason: {r['confidence_band']}."
+                )
+                continue
 
-    if isinstance(prediction, (list, tuple)):
-        prob = float(np.ravel(prediction[-1])[0])
-    else:
-        prob = float(np.ravel(prediction)[0])
+            prob_pct = float(r["probability"]) * 100.0
+            st.write(
+                f"For the uploaded file '{r['file_name']}', the model estimates a pneumonia probability of "
+                f"{prob_pct:.2f}%. This falls under '{r['confidence_band']}'. "
+                f"Based on the selected decision threshold of {threshold:.2f}, the predicted outcome is "
+                f"'{r['prediction']}'."
+            )
 
-    return max(0.0, min(1.0, prob))
+            if show_gradcam:
+                try:
+                    # Use original image for display; heatmap computed from resized input
+                    heatmap = make_gradcam_heatmap(preprocess(read_dicom(next(ff for ff in uploaded_files if ff.name == r["file_name"]))))
 
-# ----------------------------------------------------
-# Inference Section
-# ----------------------------------------------------
-st.subheader("Prediction Results")
+                    # We need original image again (Streamlit upload read pointer consumed; re-read by caching bytes)
+                    # Workaround: store bytes during first loop is better; for simplicity, skip re-read failure.
+                except Exception:
+                    pass
 
-if submit:
-    if not uploaded_files:
-        st.warning("Please upload at least one DICOM file before clicking Submit.")
-    else:
-        with st.spinner("Processing uploaded file(s)..."):
-            for file in uploaded_files:
+        # Show Grad-CAM images in a robust way (re-read bytes by caching)
+        if show_gradcam:
+            st.markdown("### Grad-CAM Heatmaps")
+            for f in uploaded_files:
                 try:
-                    image_array = read_dicom(file)
-                    x_input = preprocess(image_array)
-                    probability = get_probability(x_input)
+                    # read again safely (need cached bytes)
+                    data = f.getvalue()
+                    dcm = pydicom.dcmread(io.BytesIO(data))
+                    img = dcm.pixel_array.astype(np.float32)
+                    img = (img - img.min()) / (img.max() - img.min() + 1e-8)
 
-                    predicted_label = "Pneumonia" if probability >= threshold else "Not Pneumonia"
+                    x = preprocess(img)
+                    heatmap = make_gradcam_heatmap(x)
+                    fig = overlay_heatmap_on_image(tf.image.resize(img[..., None], (img_size, img_size)).numpy().squeeze(), heatmap)
+                    st.write(f"Heatmap for: {f.name}")
+                    st.pyplot(fig)
+                except Exception as e:
+                    st.warning(f"Could not generate Grad-CAM for {f.name}. Reason: {e}")
 
-                    st.write(
-                        f"For the uploaded file '{file.name}', the model estimates a pneumonia probability of "
-                        f"{probability * 100:.2f}%. Based on the selected decision threshold of {threshold:.2f}, "
-                        f"the predicted outcome is '{predicted_label}'."
-                    )
+        # Downloads
+        st.markdown("### Downloads")
+        csv_bytes = df.to_csv(index=False).encode("utf-8")
+        st.download_button(
+            "Download CSV",
+            data=csv_bytes,
+            file_name="predictions.csv",
+            mime="text/csv",
+            use_container_width=True
+        )
 
-                except Exception as e:
-                    st.error(
-                        f"For the uploaded file '{file.name}', the system could not generate a prediction. "
-                        f"Reason: {str(e)}."
-                    )
+        pdf_bytes = build_pdf_report(df[df["prediction"] != "Error"], threshold)
+        st.download_button(
+            "Download PDF Report",
+            data=pdf_bytes,
+            file_name="pneumonia_report.pdf",
+            mime="application/pdf",
+            use_container_width=True
+        )
 
 st.divider()
 st.caption(
-    "Clinical Notice: This application is designed for decision support purposes only. "
-    "Final diagnosis and treatment decisions must be made by qualified healthcare professionals."
+    "Clinical note: This application is designed for decision support only. Final diagnosis and treatment decisions "
+    "must be made by qualified healthcare professionals."
 )