Update src/streamlit_app.py

src/streamlit_app.py

@@ -13,25 +13,23 @@ 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",
+    page_title="Pneumonia Detection (Chest X-ray) - Clinical Decision Support",
+    layout="centered"
 )
 
-st.title("Pneumonia Detection (Chest X-ray) – Clinical Decision Support")
+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 click Submit. "
     "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")
 VERSION_PATH = os.path.join(REPO_ROOT, "model_version.json")  # optional
@@ -44,240 +42,213 @@ def load_model():
     try:
         m = keras.models.load_model(MODEL_PATH)
     except Exception:
-        # if Lambda layers / unsafe deserialization exists
+        # If you trained it, it's safe to allow deserialization
         keras.config.enable_unsafe_deserialization()
         m = keras.models.load_model(MODEL_PATH, safe_mode=False)
     return m
 
+model = load_model()
+
+# 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
+
 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")
+                return json.load(f).get("version", "ResNet50_v1")
         except Exception:
-            return "unknown"
-    return "v1"
+            return "ResNet50_v1"
+    return "ResNet50_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
-# ============================================================
+# -----------------------------
+# Confidence interpretation
+# -----------------------------
 def interpret_confidence(prob: float) -> str:
     if prob < 0.30:
         return "Low likelihood (<30%)"
     elif prob <= 0.60:
-        return "Borderline suspicion (30–60%)"
+        return "Borderline suspicion (30-60%)"
     else:
         return "High likelihood (>60%)"
 
-def read_dicom_bytes(file_bytes: bytes) -> np.ndarray:
-    dcm = pydicom.dcmread(io.BytesIO(file_bytes))
+# -----------------------------
+# DICOM + preprocessing
+# -----------------------------
+def dicom_bytes_to_img(data: bytes) -> np.ndarray:
+    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)
+    img = (img - img_min) / (img_max - img_min + 1e-8)  # 0..1
 
     return img
 
 def preprocess(img_2d: np.ndarray) -> np.ndarray:
-    """
-    (H,W) -> (1,img_size,img_size,C) float32 in 0..1
-    """
+    # (H,W) -> (1,img_size,img_size,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()  # (img_size,img_size,1)
 
     if exp_ch == 3 and x.shape[-1] == 1:
-        x = np.repeat(x, 3, axis=-1)
+        x = np.repeat(x, 3, axis=-1)     # (img_size,img_size,3)
     elif exp_ch == 1 and x.shape[-1] == 3:
-        x = x[..., :1]
+        x = x[..., :1]                  # (img_size,img_size,1)
 
-    x = np.expand_dims(x, axis=0)  # (1,img_size,img_size,C)
+    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))
 
-
-# ============================================================
-# 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, tf.keras.layers.Conv2D):
-            return layer.name
-
-    raise ValueError("No Conv2D layer found for Grad-CAM.")
-
+# -----------------------------
+# Grad-CAM (robust for nested "resnet50" submodel)
+# -----------------------------
 @st.cache_resource
-def get_gradcam_model_and_layername():
-    last_conv_name = _find_last_conv2d_layer_name(model)
+def build_gradcam_tools():
+    # If your backbone layer name is different, change here:
+    backbone_name = "resnet50"
+    backbone = model.get_layer(backbone_name)
 
-    # 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)
+    # pick the last Conv2D inside the backbone
+    last_conv = None
+    for lyr in reversed(backbone.layers):
+        if isinstance(lyr, tf.keras.layers.Conv2D):
+            last_conv = lyr.name
+            break
+    if last_conv is None:
+        raise ValueError("Could not find a Conv2D layer inside the ResNet backbone.")
 
-    grad_model = keras.Model(
-        inputs=model.inputs,
-        outputs=[conv_layer.output, model.output],
-    )
-    return grad_model, last_conv_name
+    # outputs: last conv feature map + probability head
+    conv_out = backbone.get_layer(last_conv).output
+    prob_out = model.outputs[-1] if isinstance(model.outputs, (list, tuple)) else model.output
 
-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()
-
-    with tf.GradientTape() as tape:
-        conv_outputs, preds = grad_model(img_array)
+    grad_model = keras.Model(inputs=model.inputs, outputs=[conv_out, prob_out])
+    return grad_model, last_conv
 
-        # multi-head -> take last output for probability
-        if isinstance(preds, (list, tuple)):
-            preds = preds[-1]
+def make_gradcam_heatmap(img_batch: np.ndarray) -> np.ndarray:
+    grad_model, _ = build_gradcam_tools()
 
-        # preds shape could be (1,1) or (1,)
-        loss = preds[:, 0] if preds.ndim == 2 else preds
+    x = tf.convert_to_tensor(img_batch, dtype=tf.float32)
 
-    grads = tape.gradient(loss, conv_outputs)
+    with tf.GradientTape() as tape:
+        conv_out, preds = grad_model(x, training=False)
+        loss = preds[:, 0]  # binary sigmoid prob
 
-    # safety in case grads is None
+    grads = tape.gradient(loss, conv_out)
     if grads is None:
-        raise ValueError("Gradients are None. Grad-CAM cannot be computed for this model output.")
+        raise ValueError("Gradients are None. Grad-CAM cannot be computed for this model setup.")
 
-    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))  # (channels,)
-    conv_outputs = conv_outputs[0]  # (Hc,Wc,channels)
+    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
+    conv_out = conv_out[0]  # (H,W,channels)
 
-    heatmap = tf.reduce_sum(conv_outputs * pooled_grads, axis=-1)  # (Hc,Wc)
+    heatmap = tf.reduce_sum(conv_out * pooled_grads, axis=-1)
     heatmap = tf.maximum(heatmap, 0)
 
-    denom = tf.reduce_max(heatmap)
-    heatmap = heatmap / (denom + 1e-8)
-
+    denom = tf.reduce_max(heatmap) + 1e-8
+    heatmap = heatmap / denom
     return heatmap.numpy()
 
-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
-    """
+def overlay_heatmap(img_2d: np.ndarray, heatmap: np.ndarray):
+    # resize heatmap to img_size
     heat = tf.image.resize(heatmap[..., None], (img_size, img_size)).numpy().squeeze()
 
+    # ensure base image displayed at img_size
+    base = tf.image.resize(img_2d[..., None], (img_size, img_size)).numpy().squeeze()
+
     fig = plt.figure(figsize=(5, 5))
-    plt.imshow(img_2d_resized, cmap="gray")
+    plt.imshow(base, cmap="gray")
     plt.imshow(heat, cmap="jet", alpha=0.35)
     plt.axis("off")
     plt.tight_layout()
     return fig
 
+# -----------------------------
+# PDF generation (fix unicode issues)
+# -----------------------------
+def safe_text(s: str, max_len: int = 180) -> str:
+    """
+    Convert to something FPDF core fonts can render.
+    Also trims very long strings to avoid layout failures.
+    """
+    if s is None:
+        return ""
+    s = str(s)
+
+    # Replace common unicode dashes/quotes with ascii
+    s = s.replace("–", "-").replace("—", "-").replace("’", "'").replace("“", '"').replace("”", '"')
+
+    # Remove / replace any remaining unsupported chars (latin-1 fallback)
+    s = s.encode("latin-1", "replace").decode("latin-1")
 
-# ============================================================
-# PDF generator (stable)
-# ============================================================
-class PDF(FPDF):
-    pass
+    # Avoid extremely long unbroken lines
+    if len(s) > max_len:
+        s = s[:max_len] + "..."
+    return s
 
-def build_pdf_report(df: pd.DataFrame, threshold: float, model_version: str) -> bytes:
-    pdf = PDF()
+def build_pdf_report(df_ok: pd.DataFrame, threshold: float, model_version: str) -> bytes:
+    pdf = FPDF()
     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.set_font("Arial", size=10)
-
-    # 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.set_font("Helvetica", size=12)
 
-        pdf.ln(2)
+    pdf.cell(0, 8, safe_text("Pneumonia Detection Report"), ln=True)
+    pdf.set_font("Helvetica", size=10)
+    pdf.cell(0, 6, safe_text(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"), ln=True)
+    pdf.cell(0, 6, safe_text(f"Model Version: {model_version}"), ln=True)
+    pdf.cell(0, 6, safe_text(f"Decision Threshold: {threshold:.2f}"), ln=True)
+    pdf.ln(4)
 
-    return pdf.output(dest="S").encode("latin-1")
+    pdf.set_font("Helvetica", size=10)
+    for _, row in df_ok.iterrows():
+        line1 = f"File: {row['file_name']}"
+        line2 = f"Probability: {row['probability']*100:.2f}%"
+        line3 = f"Confidence: {row['confidence_level']}"
+        line4 = f"Prediction: {row['prediction']}"
 
+        for line in [line1, line2, line3, line4]:
+            pdf.multi_cell(0, 6, safe_text(line))
 
-# ============================================================
+        pdf.ln(2)
+
+    out = pdf.output(dest="S")
+    # fpdf may return str in some versions
+    if isinstance(out, str):
+        out = out.encode("latin-1", "ignore")
+    return out
+
+# -----------------------------
 # UI
-# ============================================================
+# -----------------------------
 st.subheader("Model Parameters")
 
 threshold = st.slider(
     "Decision Threshold",
     min_value=0.01,
     max_value=0.99,
-    value=0.37,  # ResNet default (your best thr)
+    value=0.37,   # your ResNet best threshold default
     step=0.01,
-    help="If predicted probability ≥ threshold → Pneumonia, else → Not Pneumonia.",
+    help="If predicted probability is greater than or equal to the threshold, output is Pneumonia. Otherwise Not Pneumonia."
 )
 
 show_gradcam = st.checkbox("Show Grad-CAM heatmap (explainability)", value=True)
 
 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)
@@ -291,54 +262,44 @@ 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)
+        # cache bytes once (so we can read multiple times safely)
+        file_bytes = {f.name: f.getvalue() for f in uploaded_files}
 
+        rows = []
         with st.spinner("Running inference..."):
-            for f in uploaded_files:
+            for name, data in file_bytes.items():
                 ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                 try:
-                    b = f.getvalue()
-                    img = read_dicom_bytes(b)  # (H,W) 0..1
+                    img = dicom_bytes_to_img(data)
                     x = preprocess(img)
                     prob = predict_prob(x)
 
                     pred_label = "Pneumonia" if prob >= threshold else "Not Pneumonia"
                     conf_level = interpret_confidence(prob)
 
-                    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))
-
+                    rows.append({
+                        "timestamp": ts,
+                        "model_version": MODEL_VERSION,
+                        "file_name": name,
+                        "probability": prob,
+                        "prediction": pred_label,
+                        "confidence_level": conf_level,
+                        "error": ""
+                    })
                 except Exception as e:
-                    rows.append(
-                        {
-                            "timestamp": ts,
-                            "model_version": MODEL_VERSION,
-                            "file_name": f.name,
-                            "probability": np.nan,
-                            "confidence_level": "NA",
-                            "prediction": "Error",
-                            "error": str(e),
-                        }
-                    )
+                    rows.append({
+                        "timestamp": ts,
+                        "model_version": MODEL_VERSION,
+                        "file_name": name,
+                        "probability": np.nan,
+                        "prediction": "Error",
+                        "confidence_level": "",
+                        "error": safe_text(str(e), max_len=140)
+                    })
 
         df = pd.DataFrame(rows)
 
@@ -347,34 +308,34 @@ 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.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']}'."
+                    f"Reason: {r['error']}."
                 )
+                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_level']}'. "
+                f"Based on the selected decision threshold of {threshold:.2f}, the predicted outcome is "
+                f"'{r['prediction']}'."
+            )
 
         # Grad-CAM section
         if show_gradcam:
             st.markdown("### Grad-CAM Heatmaps")
-
-            # Resize original image for display
-            for (fname, _, img_2d) in file_cache:
+            for name, data in file_bytes.items():
                 try:
-                    img_resized = tf.image.resize(img_2d[..., None], (img_size, img_size)).numpy().squeeze()
-                    x = preprocess(img_2d)
+                    img = dicom_bytes_to_img(data)
+                    x = preprocess(img)
                     heatmap = make_gradcam_heatmap(x)
-                    fig = overlay_heatmap_on_image(img_resized, heatmap)
-                    st.write(f"Heatmap for: {fname}")
+                    fig = overlay_heatmap(img, heatmap)
+                    st.write(f"Heatmap for: {name}")
                     st.pyplot(fig)
                 except Exception as 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}")
+                    st.warning(
+                        f"Could not generate Grad-CAM for {name}. "
+                        f"Reason: {safe_text(str(e), max_len=160)}"
+                    )
 
         # Downloads
         st.markdown("### Downloads")
@@ -385,7 +346,7 @@ if submit:
             data=csv_bytes,
             file_name="predictions.csv",
             mime="text/csv",
-            use_container_width=True,
+            use_container_width=True
         )
 
         df_ok = df[df["prediction"] != "Error"].copy()
@@ -396,11 +357,10 @@ if submit:
                 data=pdf_bytes,
                 file_name="pneumonia_report.pdf",
                 mime="application/pdf",
-                use_container_width=True,
+                use_container_width=True
             )
         else:
-            st.info("PDF report is available once at least one file is successfully processed.")
-
+            st.info("PDF report is available only when at least one file is successfully processed.")
 
 st.divider()
 st.caption(