Update src/streamlit_app.py

src/streamlit_app.py

@@ -13,12 +13,13 @@ import matplotlib.pyplot as plt
 
 from fpdf import FPDF
 
-# -----------------------------
+
+# ============================================================
 # Page config
-# -----------------------------
+# ============================================================
 st.set_page_config(
     page_title="Pneumonia Detection (Chest X-ray) – Clinical Decision Support",
-    layout="centered"
+    layout="centered",
 )
 
 st.title("Pneumonia Detection (Chest X-ray) – Clinical Decision Support")
@@ -27,14 +28,13 @@ st.caption(
     "This tool is for decision support only and does not replace clinical judgment."
 )
 
-# -----------------------------
+
+# ============================================================
 # 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")
+VERSION_PATH = os.path.join(REPO_ROOT, "model_version.json")  # optional
 
 @st.cache_resource
 def load_model():
@@ -44,20 +44,11 @@ def load_model():
     try:
         m = keras.models.load_model(MODEL_PATH)
     except Exception:
+        # if Lambda layers / unsafe deserialization exists
         keras.config.enable_unsafe_deserialization()
         m = keras.models.load_model(MODEL_PATH, safe_mode=False)
     return m
 
-model = load_model()
-
-# 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
-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:
@@ -68,10 +59,27 @@ def get_model_version():
     return "v1"
 
 MODEL_VERSION = get_model_version()
+model = load_model()
+
+# model input details: (None, H, W, C)
+input_shape = model.input_shape
+img_size = int(input_shape[1]) if input_shape and input_shape[1] else 256
+exp_ch = int(input_shape[-1]) if input_shape and input_shape[-1] else 1
+
+
+# ============================================================
+# Helpers
+# ============================================================
+def interpret_confidence(prob: float) -> str:
+    if prob < 0.30:
+        return "Low likelihood (<30%)"
+    elif prob <= 0.60:
+        return "Borderline suspicion (30–60%)"
+    else:
+        return "High likelihood (>60%)"
 
-def read_dicom(uploaded_file) -> np.ndarray:
-    data = uploaded_file.read()
-    dcm = pydicom.dcmread(io.BytesIO(data))
+def read_dicom_bytes(file_bytes: bytes) -> np.ndarray:
+    dcm = pydicom.dcmread(io.BytesIO(file_bytes))
     img = dcm.pixel_array.astype(np.float32)
 
     # Normalize to 0..1
@@ -82,143 +90,184 @@ def read_dicom(uploaded_file) -> np.ndarray:
     return img
 
 def preprocess(img_2d: np.ndarray) -> np.ndarray:
-    # (H,W) -> (1,H,W,C) float32 0..1
+    """
+    (H,W) -> (1,img_size,img_size,C) float32 in 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()
+    x = x.numpy()  # (img_size,img_size,1)
 
     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)
+    x = np.expand_dims(x, axis=0)  # (1,img_size,img_size,C)
     return x.astype(np.float32)
 
 def predict_prob(x: np.ndarray) -> float:
+    """
+    Supports single-head and multi-head models.
+    Uses last output as probability when outputs are list/tuple.
+    """
     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
+# ============================================================
+# Grad-CAM (robust layer selection)
+# ============================================================
+def _find_last_conv2d_layer_name(m: keras.Model) -> str:
+    # Prefer backbone conv layers if a common backbone layer exists
+    backbone_names = ["resnet50", "ResNet50", "backbone"]
+    for nm in backbone_names:
+        try:
+            bb = m.get_layer(nm)
+            # walk backwards in backbone
+            for layer in reversed(bb.layers):
+                if isinstance(layer, tf.keras.layers.Conv2D):
+                    return f"{nm}/{layer.name}"
+        except Exception:
+            pass
+
+    # Fallback: scan the whole model
     for layer in reversed(m.layers):
-        if isinstance(layer, keras.layers.Conv2D):
+        if isinstance(layer, tf.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.")
+
+    raise ValueError("No Conv2D layer found 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 get_gradcam_model_and_layername():
+    last_conv_name = _find_last_conv2d_layer_name(model)
 
-def make_gradcam_heatmap(x_input: np.ndarray) -> np.ndarray:
-    grad_model, _ = get_gradcam_model(model)
+    # If name includes "backbone/layer", resolve it
+    if "/" in last_conv_name:
+        parent, child = last_conv_name.split("/", 1)
+        conv_layer = model.get_layer(parent).get_layer(child)
+    else:
+        conv_layer = model.get_layer(last_conv_name)
+
+    grad_model = keras.Model(
+        inputs=model.inputs,
+        outputs=[conv_layer.output, model.output],
+    )
+    return grad_model, last_conv_name
+
+def make_gradcam_heatmap(img_array: np.ndarray) -> np.ndarray:
+    """
+    img_array: (1,H,W,C)
+    returns heatmap: (Hc, Wc) normalized 0..1
+    """
+    grad_model, _ = get_gradcam_model_and_layername()
 
-    x_tensor = tf.convert_to_tensor(x_input, dtype=tf.float32)
     with tf.GradientTape() as tape:
-        conv_out, preds = grad_model(x_tensor)
+        conv_outputs, preds = grad_model(img_array)
 
+        # multi-head -> take last output for probability
         if isinstance(preds, (list, tuple)):
             preds = preds[-1]
 
-        # binary prob is preds[:,0]
-        score = preds[:, 0]
+        # preds shape could be (1,1) or (1,)
+        loss = preds[:, 0] if preds.ndim == 2 else preds
+
+    grads = tape.gradient(loss, conv_outputs)
 
-    grads = tape.gradient(score, conv_out)
-    pooled = tf.reduce_mean(grads, axis=(0, 1, 2))
-    conv_out = conv_out[0]
+    # safety in case grads is None
+    if grads is None:
+        raise ValueError("Gradients are None. Grad-CAM cannot be computed for this model output.")
 
-    heatmap = conv_out @ pooled[..., tf.newaxis]
-    heatmap = tf.squeeze(heatmap)
+    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))  # (channels,)
+    conv_outputs = conv_outputs[0]  # (Hc,Wc,channels)
 
+    heatmap = tf.reduce_sum(conv_outputs * pooled_grads, axis=-1)  # (Hc,Wc)
     heatmap = tf.maximum(heatmap, 0)
-    denom = tf.reduce_max(heatmap) + 1e-8
-    heatmap = heatmap / denom
+
+    denom = tf.reduce_max(heatmap)
+    heatmap = heatmap / (denom + 1e-8)
+
     return heatmap.numpy()
 
-def overlay_heatmap_on_image(img_2d: np.ndarray, heatmap: np.ndarray):
-    # Resize heatmap to img_size
+def overlay_heatmap_on_image(img_2d_resized: np.ndarray, heatmap: np.ndarray):
+    """
+    img_2d_resized: (img_size,img_size) in 0..1
+    heatmap: (Hc,Wc) in 0..1
+    """
     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(img_2d_resized, 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 generator (stable)
+# ============================================================
+class PDF(FPDF):
+    pass
+
+def build_pdf_report(df: pd.DataFrame, threshold: float, model_version: str) -> bytes:
+    pdf = PDF()
+    pdf.set_auto_page_break(auto=True, margin=12)
     pdf.add_page()
+    pdf.set_margins(12, 12, 12)
+
     pdf.set_font("Arial", size=12)
+    pdf.cell(0, 8, "Pneumonia Detection Report", ln=True)
+
+    pdf.set_font("Arial", size=10)
+    pdf.cell(0, 7, f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", ln=True)
+    pdf.cell(0, 7, f"Model Version: {model_version}", ln=True)
+    pdf.cell(0, 7, f"Decision Threshold: {threshold:.2f}", ln=True)
+    pdf.ln(3)
 
-    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()
+
+    # Write each row safely
+    for _, row in df.iterrows():
+        pdf.set_x(pdf.l_margin)
+        prob = row.get("probability", np.nan)
+        prob_pct = "NA" if pd.isna(prob) else f"{float(prob)*100:.2f}%"
+
+        lines = [
+            f"File: {row.get('file_name','')}",
+            f"Probability: {prob_pct}",
+            f"Confidence Level: {row.get('confidence_level','')}",
+            f"Prediction: {row.get('prediction','')}",
+            f"Timestamp: {row.get('timestamp','')}",
+        ]
+
+        for line in lines:
+            # Use multi_cell for wrapping and reset x each time
+            pdf.set_x(pdf.l_margin)
+            pdf.multi_cell(0, 6, line)
+
+        pdf.ln(2)
 
     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,   # your ResNet best-thr default
+    value=0.37,  # ResNet default (your best thr)
     step=0.01,
-    help="If predicted probability ≥ threshold → Pneumonia, else → Not Pneumonia."
+    help="If predicted probability ≥ threshold → Pneumonia, else → Not Pneumonia.",
 )
 
 show_gradcam = st.checkbox("Show Grad-CAM heatmap (explainability)", value=True)
@@ -228,7 +277,7 @@ st.subheader("Upload Chest X-ray DICOM Files")
 uploaded_files = st.file_uploader(
     "Select one or multiple DICOM files (.dcm)",
     type=["dcm"],
-    accept_multiple_files=True
+    accept_multiple_files=True,
 )
 
 col1, col2 = st.columns(2)
@@ -242,38 +291,54 @@ if clear:
 
 st.subheader("Prediction Results")
 
+
+# ============================================================
+# Inference
+# ============================================================
 if submit:
     if not uploaded_files:
         st.warning("Please upload at least one DICOM file before submitting.")
     else:
         rows = []
+        file_cache = []  # (filename, bytes, img_2d_norm)
+
         with st.spinner("Running inference..."):
             for f in uploaded_files:
+                ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                 try:
-                    img = read_dicom(f)
+                    b = f.getvalue()
+                    img = read_dicom_bytes(b)  # (H,W) 0..1
                     x = preprocess(img)
                     prob = predict_prob(x)
+
                     pred_label = "Pneumonia" if prob >= threshold else "Not Pneumonia"
-                    band = confidence_bucket(prob)
+                    conf_level = interpret_confidence(prob)
 
-                    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
-                    })
+                    rows.append(
+                        {
+                            "timestamp": ts,
+                            "model_version": MODEL_VERSION,
+                            "file_name": f.name,
+                            "probability": prob,
+                            "confidence_level": conf_level,
+                            "prediction": pred_label,
+                        }
+                    )
+
+                    file_cache.append((f.name, b, img))
 
                 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)
-                    })
+                    rows.append(
+                        {
+                            "timestamp": ts,
+                            "model_version": MODEL_VERSION,
+                            "file_name": f.name,
+                            "probability": np.nan,
+                            "confidence_level": "NA",
+                            "prediction": "Error",
+                            "error": str(e),
+                        }
+                    )
 
         df = pd.DataFrame(rows)
 
@@ -282,66 +347,60 @@ if submit:
             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']}."
+                    f"Reason: {r.get('error','Unknown error')}."
+                )
+            else:
+                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_level']}'. "
+                    f"Based on the selected decision threshold of {threshold:.2f}, the predicted outcome is "
+                    f"'{r['prediction']}'."
                 )
-                continue
-
-            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']}'."
-            )
-
-            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"]))))
-
-                    # 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
 
-        # Show Grad-CAM images in a robust way (re-read bytes by caching)
+        # Grad-CAM section
         if show_gradcam:
             st.markdown("### Grad-CAM Heatmaps")
-            for f in uploaded_files:
-                try:
-                    # 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)
 
-                    x = preprocess(img)
+            # Resize original image for display
+            for (fname, _, img_2d) in file_cache:
+                try:
+                    img_resized = tf.image.resize(img_2d[..., None], (img_size, img_size)).numpy().squeeze()
+                    x = preprocess(img_2d)
                     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}")
+                    fig = overlay_heatmap_on_image(img_resized, heatmap)
+                    st.write(f"Heatmap for: {fname}")
                     st.pyplot(fig)
                 except Exception as e:
-                    st.warning(f"Could not generate Grad-CAM for {f.name}. Reason: {e}")
+                    # This will now show the *actual* layer list + reason,
+                    # instead of failing with a wrong hard-coded layer name.
+                    st.warning(f"Could not generate Grad-CAM for {fname}. Reason: {e}")
 
         # 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
+            use_container_width=True,
         )
 
-        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
-        )
+        df_ok = df[df["prediction"] != "Error"].copy()
+        if len(df_ok) > 0:
+            pdf_bytes = build_pdf_report(df_ok, threshold, MODEL_VERSION)
+            st.download_button(
+                "Download PDF Report",
+                data=pdf_bytes,
+                file_name="pneumonia_report.pdf",
+                mime="application/pdf",
+                use_container_width=True,
+            )
+        else:
+            st.info("PDF report is available once at least one file is successfully processed.")
+
 
 st.divider()
 st.caption(