Update streamlit_app.py

streamlit_app.py

@@ -1,151 +1,140 @@
 # streamlit_app.py
-import io
-import os
-import numpy as np
-from PIL import Image
 import streamlit as st
+import numpy as np
 import tensorflow as tf
-from tensorflow.keras.models import load_model 
-from tensorflow.keras.applications.densenet import preprocess_input as densenet_preprocess
+from tensorflow.keras.models import load_model, Model
+import cv2
 import pydicom
-from pydicom.pixel_data_handlers.util import apply_voi_lut
-import matplotlib.cm as cm
-
-# -------- CONFIG --------
-MODEL_FILENAME = "Model2_exact_serialized.keras"   # model file expected in app folder
-IMG_SIZE = (224, 224)
-THRESHOLD = 0.62
-ENABLE_GRADCAM = True
-# ------------------------
+from PIL import Image
 
-st.set_page_config(page_title="Pneumonia Detection (CheXNet)", layout="centered")
+st.set_page_config(page_title="Pneumonia Detection", layout="wide")
 
-st.title("Pneumonia detection (CheXNet)")
-st.write("Upload a chest X-ray (DICOM or PNG/JPG). The app predicts probability of pneumonia.")
 
-# ------- utilities -------
-def dicom_to_image_array(dicom_bytes):
-    ds = pydicom.dcmread(io.BytesIO(dicom_bytes), force=True)
-    try:
-        arr = ds.pixel_array
-    except Exception as e:
-        raise RuntimeError(f"Could not decode DICOM pixel data: {e}")
-    if arr.ndim == 3:
-        arr = arr[0]
-    try:
-        arr = apply_voi_lut(arr, ds)
-    except Exception:
-        pass
-    arr = arr.astype(np.float32)
-    if getattr(ds, "PhotometricInterpretation", "").upper() == "MONOCHROME1":
-        arr = np.max(arr) - arr
-    mn, mx = arr.min(), arr.max()
-    if mx > mn:
-        arr = (arr - mn) / (mx - mn)
-    else:
-        arr = arr - mn
-    arr = (arr * 255.0).clip(0,255).astype(np.uint8)
-    return arr
-
-def to_rgb_uint8_from_upload(uploaded_file):
-    """Return RGB uint8 (H,W,3) array resized to IMG_SIZE."""
-    if uploaded_file is None:
-        raise RuntimeError("No file")
-    raw = uploaded_file.read()
-    # try DICOM
-    try:
-        ds = pydicom.dcmread(io.BytesIO(raw), stop_before_pixels=True, force=True)
-        if hasattr(ds, "PixelData") or getattr(ds, "Rows", None):
-            arr = dicom_to_image_array(raw)
-            if arr.ndim == 2:
-                arr = np.stack([arr]*3, axis=-1)
-            pil = Image.fromarray(arr).convert("RGB").resize(IMG_SIZE)
-            return np.array(pil)
-    except Exception:
-        pass
-    # fallback normal image
-    try:
-        pil = Image.open(io.BytesIO(raw)).convert("L").resize(IMG_SIZE)
-        arr = np.stack([np.array(pil)]*3, axis=-1)
-        return arr.astype(np.uint8)
-    except Exception as e:
-        raise RuntimeError("Unsupported file format. Upload a DICOM or PNG/JPG.") from e
-
-# -------- model load (cached) --------
+# ---------------------------
+# Load Model
+# ---------------------------
 @st.cache_resource
-def load_predict_model(model_path):
-    if not os.path.exists(model_path):
-        raise FileNotFoundError(f"Model file not found: {model_path}")
-    m = load_model(model_path, compile=False)
-    return m
-
-# Grad-CAM utilities
-def find_last_conv_layer(m):
-    for layer in reversed(m.layers):
-        out_shape = getattr(layer, "output_shape", None)
-        if out_shape and len(out_shape) == 4 and "conv" in layer.name:
-            return layer.name
-    return m.layers[-3].name
-
-def make_gradcam_image(rgb_uint8, model, last_conv_name=None, alpha=0.4, cmap_name="jet"):
-    img = rgb_uint8.astype(np.float32)
-    if last_conv_name is None:
-        last_conv_name = find_last_conv_layer(model)
-    grad_model = tf.keras.models.Model([model.inputs], [model.get_layer(last_conv_name).output, model.output])
-    x = densenet_preprocess(np.expand_dims(img.astype(np.float32), axis=0))
+def load_my_model():
+    model = load_model("model/best_model.keras", compile=False)
+    return model
+
+
+model = load_my_model()
+
+
+# ---------------------------
+# Preprocess image
+# ---------------------------
+def load_image(file):
+    """Loads PNG/JPG/DICOM and returns a grayscale 224x224 normalized array."""
+    filename = file.name.lower()
+
+    if filename.endswith(".dcm"):
+        dcm = pydicom.dcmread(file)
+        img = dcm.pixel_array.astype(np.float32)
+        img = cv2.resize(img, (224, 224))
+        img = img / np.max(img)
+        return img
+
+    # PNG / JPG / JPEG
+    img = Image.open(file).convert("L")
+    img = img.resize((224, 224))
+    img = np.array(img).astype(np.float32) / 255.0
+    return img
+
+
+# ---------------------------
+# Robust Grad-CAM
+# ---------------------------
+def grad_cam(model, img_array, layer_name=None, eps=1e-8):
+    """
+    img_array: (224,224) normalized grayscale → will be expanded to (1,224,224,1)
+    """
+
+    # Expand dims for model
+    x = np.expand_dims(img_array, axis=0)          # (1,224,224)
+    x = np.expand_dims(x, axis=-1)                 # (1,224,224,1)
+    x = tf.convert_to_tensor(x, dtype=tf.float32)
+
+    # Auto-detect last conv layer if not provided
+    if layer_name is None:
+        for layer in reversed(model.layers):
+            if hasattr(layer, "output_shape") and len(layer.output_shape) == 4:
+                layer_name = layer.name
+                break
+
+    last_conv = model.get_layer(layer_name)
+    grad_model = Model([model.inputs], [last_conv.output, model.output])
+
     with tf.GradientTape() as tape:
-        conv_outputs, preds = grad_model(x)
-        loss = preds[:, 0]
+        conv_outputs, predictions = grad_model(x)
+        class_idx = int(tf.argmax(predictions[0]))
+        loss = predictions[:, class_idx]
+
     grads = tape.gradient(loss, conv_outputs)
-    weights = tf.reduce_mean(grads, axis=(1,2))
-    cam = tf.reduce_sum(tf.multiply(weights[:, tf.newaxis, tf.newaxis, :], conv_outputs), axis=-1)
-    cam = tf.squeeze(cam).numpy()
+    if grads is None:
+        raise RuntimeError("Gradients are None. Model may not be connected properly.")
+
+    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2)).numpy()
+
+    conv_outputs = conv_outputs[0].numpy()  # (H,W,channels)
+
+    for i in range(len(pooled_grads)):
+        conv_outputs[:, :, i] *= pooled_grads[i]
+
+    cam = np.mean(conv_outputs, axis=-1)
     cam = np.maximum(cam, 0)
-    cam_max = cam.max() if cam.max() != 0 else 1e-8
-    cam = cam / cam_max
-    cam_img = Image.fromarray(np.uint8(cam * 255)).resize((img.shape[1], img.shape[0]), resample=Image.BILINEAR)
-    cam_arr = np.array(cam_img).astype(np.float32)/255.0
-    colormap = cm.get_cmap(cmap_name)
-    heatmap = colormap(cam_arr)[:, :, :3]
-    heat_uint8 = np.uint8(heatmap * 255)
-    heat_pil = Image.fromarray(heat_uint8).convert("RGBA").resize((img.shape[1], img.shape[0]))
-    base_pil = Image.fromarray(np.uint8(img)).convert("RGBA")
-    blended = Image.blend(base_pil, heat_pil, alpha=alpha)
-    return blended.convert("RGB")
-
-# -------- UI elements --------
-col1, col2 = st.columns([1,1])
-with col1:
-    uploaded = st.file_uploader("Upload DICOM or PNG/JPG", type=["dcm","png","jpg","jpeg","tif","tiff"])
-with col2:
-    thresh = st.number_input("Decision threshold (probability)", min_value=0.0, max_value=1.0, value=float(THRESHOLD), step=0.01)
-
-if uploaded is not None:
+
+    cam -= cam.min()
+    if cam.max() > eps:
+        cam /= cam.max() + eps
+
+    cam = cv2.resize(cam, (224, 224))
+    return cam
+
+
+# ---------------------------
+# UI
+# ---------------------------
+st.title("🫁 Pneumonia Detection")
+st.write("Upload a chest scan image (DICOM or PNG/JPG).")
+
+
+file = st.file_uploader("Upload Image", type=["png", "jpg", "jpeg", "dcm"])
+
+if file:
     try:
-        rgb = to_rgb_uint8_from_upload(uploaded)
-    except Exception as e:
-        st.error(f"Failed to process file: {e}")
-        st.stop()
+        img = load_image(file)
 
-    st.image(rgb, caption="Input (resized)", use_container_width=True)
+        st.subheader("Input Image")
+        st.image(img, caption="Uploaded Image", use_container_width=True, clamp=True)
 
-    # load model (cached)
-    model = load_predict_model(MODEL_FILENAME)
+        # Model input format (1,224,224,1)
+        x = np.expand_dims(img, axis=(0, -1))
 
-    # predict
-    x_pre = densenet_preprocess(np.expand_dims(rgb.astype(np.float32), axis=0))
-    prob = float(model.predict(x_pre, verbose=0).ravel()[0])
-    pred = int(prob >= thresh)
+        pred = model.predict(x)[0][0]
+        label = "Pneumonia" if pred >= 0.5 else "Normal"
 
-    st.markdown(f"**Pneumonia probability:** `{prob:.4f}`")
-    st.markdown(f"**Predicted class (binary):** `{pred}`  — **{'Pneumonia' if pred==1 else 'Normal'}**")
+        st.subheader("Prediction")
+        st.write(f"**Class:** {label}")
+        st.write(f"**Probability:** {float(pred):.4f}")
 
-    if ENABLE_GRADCAM:
+        # Grad-CAM
+        st.subheader("Grad-CAM Heatmap")
         try:
-            cam = make_gradcam_image(rgb, model)
-            st.image(cam, caption="Grad-CAM overlay", use_container_width=True)
+            cam = grad_cam(model, img)
+            heatmap = cv2.applyColorMap(
+                np.uint8(255 * cam), cv2.COLORMAP_JET
+            )
+            heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
+
+            overlay = 0.4 * heatmap + 0.6 * np.stack([img*255]*3, axis=-1)
+            overlay = overlay.astype(np.uint8)
+
+            st.image(overlay, caption="Grad-CAM", use_container_width=True)
+
         except Exception as e:
-            st.warning(f"Grad-CAM failed: {e}")
+            st.error(f"Grad-CAM failed: {e}")
 
-else:
-    st.info("Upload a DICOM or PNG/JPG image to run inference.")
+    except Exception as e:
+        st.error(f"Error loading image: {e}")