Update app.py

app.py

@@ -8,6 +8,7 @@ from keras.models import Model
 from keras.saving import register_keras_serializable
 from keras.layers import TFSMLayer
 import matplotlib.pyplot as plt
+import matplotlib.cm as cm
 from lime import lime_image
 from skimage.segmentation import mark_boundaries
 
@@ -40,28 +41,28 @@ class CustomTFOpLambda(tf.keras.layers.Layer):
         config.update({"function": self.function})
         return config
 
-# --- Constants ---
+# Constants
 IMG_SIZE = (224, 224)
 CLASS_NAMES = ['Normal', 'Diabetes', 'Glaucoma', 'Cataract', 'AMD', 'Hypertension', 'Myopia', 'Others']
 
-# --- Load SavedModel as TFSMLayer wrapped model ---
+# Load model using TFSMLayer
 @st.cache_resource
 def load_model():
-    model_path = "Model"  # Your SavedModel directory path
+    model_path = "Model"  # Adjust path to your SavedModel directory
     if not os.path.exists(model_path):
         st.error(f"❌ Model directory '{model_path}' not found!")
         st.stop()
     try:
         tfsm_layer = TFSMLayer(model_path, call_endpoint="serving_default")
-        inputs = Input(shape=(224, 224, 3))
+        inputs = Input(shape=(IMG_SIZE[0], IMG_SIZE[1], 3))
         outputs = tfsm_layer(inputs)
         model = Model(inputs=inputs, outputs=outputs)
         return model
     except Exception as e:
-        st.error(f"❌ Error loading model with TFSMLayer: {str(e)}")
+        st.error(f"❌ Error loading model: {str(e)}")
         st.stop()
 
-# --- Preprocessing functions ---
+# Preprocessing functions
 def crop_circle(img):
     h, w = img.shape[:2]
     center = (w // 2, h // 2)
@@ -92,93 +93,113 @@ def preprocess_image(img):
     clahe = apply_clahe(circ)
     sharp = sharpen_image(clahe)
     resized = resize_normalize(sharp)
-    return circ, clahe, sharp, resized
+    return resized
 
-# --- Grad-CAM ---
-def show_gradcam(model, img, class_idx):
-    last_conv_layer = None
+# --- Find Last Conv Layer for Grad-CAM ---
+def find_last_conv_layer(model):
     for layer in reversed(model.layers):
-        if isinstance(layer, tf.keras.layers.Conv2D):
-            last_conv_layer = layer.name
-            break
-    if last_conv_layer is None:
-        st.warning("⚠️ No Conv2D layer found for Grad-CAM.")
-        return
-    
-    grad_model = Model(model.inputs, [model.get_layer(last_conv_layer).output, model.output])
-    img_tensor = tf.convert_to_tensor(img[np.newaxis, ...])
-    
+        # Check if Conv2D or custom attention output layer name pattern
+        if isinstance(layer, tf.keras.layers.Conv2D) or 'mhsa_output' in layer.name:
+            return layer.name
+    raise ValueError("No suitable conv layer found.")
+
+# Grad-CAM generation
+def generate_gradcam(model, img_array, class_index, layer_name):
+    grad_model = tf.keras.models.Model([model.inputs], [model.get_layer(layer_name).output, model.output])
     with tf.GradientTape() as tape:
-        conv_outputs, predictions = grad_model(img_tensor)
-        loss = predictions[:, class_idx]
-    grads = tape.gradient(loss, conv_outputs)[0]
-    cam = tf.reduce_mean(grads, axis=-1).numpy()
-    
-    cam = np.maximum(cam, 0)
-    cam = cv2.resize(cam, IMG_SIZE)
-    cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-10)
-    
-    heatmap = np.uint8(255 * cam)
-    heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
-    overlay = cv2.addWeighted(np.uint8(img * 255), 0.6, heatmap, 0.4, 0)
-    
-    st.subheader("🔥 Grad-CAM")
-    st.image(overlay, use_container_width=True)
-
-# --- LIME ---
-def show_lime(model, img, class_idx):
-    explainer = lime_image.LimeImageExplainer()
-    
-    def predict_fn(images):
-        images = np.array(images)
-        preds = model.predict(images)
-        if isinstance(preds, dict):
-            preds = list(preds.values())[0]
-        return preds
-    
-    explanation = explainer.explain_instance(np.uint8(img * 255), predict_fn, top_labels=1, hide_color=0, num_samples=1000)
-    lime_img, mask = explanation.get_image_and_mask(class_idx, positive_only=True, hide_rest=False)
-    
-    st.subheader("🟢 LIME Explanation")
-    st.image(mark_boundaries(lime_img, mask), use_container_width=True)
-
-# --- Streamlit UI ---
-st.set_page_config(page_title="🧠 Retina Disease Classifier", layout="centered")
-st.title("🧠 Retina Disease Classifier")
+        conv_outputs, predictions = grad_model(img_array)
+        loss = predictions[:, class_index]
+    grads = tape.gradient(loss, conv_outputs)
+    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
+    heatmap = conv_outputs[0] @ pooled_grads[..., tf.newaxis]
+    heatmap = tf.squeeze(heatmap)
+    heatmap = tf.maximum(heatmap, 0) / (tf.math.reduce_max(heatmap) + 1e-10)
+    return heatmap.numpy()
+
+# LIME explainer
+explainer = lime_image.LimeImageExplainer()
+def predict_fn(images):
+    images = np.array(images)
+    preds = model.predict(images, verbose=0)
+    # If output is dict (due to TFSMLayer), extract predictions properly
+    if isinstance(preds, dict):
+        preds = list(preds.values())[0]
+    return preds
+
+# Explanation text per class
+explanation_text = {
+    'Normal': "Model predicted Normal based on healthy optic disc and macula.",
+    'Diabetes': "Detected retinal blood vessel changes suggestive of Diabetes.",
+    'Glaucoma': "Detected increased cupping in the optic disc indicating Glaucoma.",
+    'Cataract': "Image blur indicated potential Cataract.",
+    'AMD': "Degeneration signs in macula indicate AMD.",
+    'Hypertension': "Blood vessel narrowing/hemorrhages indicate Hypertension.",
+    'Myopia': "Tilted disc and fundus shape suggest Myopia.",
+    'Others': "Non-specific features detected, marked as Others."
+}
+
+# Visualization in Streamlit
+def display_combined_visualization(img, true_label, pred_label, pred_idx, layer_name):
+    input_array = np.expand_dims(img, axis=0)
+
+    # Grad-CAM heatmap
+    heatmap = generate_gradcam(model, input_array, pred_idx, layer_name)
+    heatmap = cv2.resize(heatmap, IMG_SIZE)
+    heatmap = np.uint8(255 * heatmap)
+    heatmap = cv2.GaussianBlur(heatmap, (7, 7), 0)
+    heatmap_rgb = cm.jet(heatmap / 255.0)[..., :3]
+    heatmap_rgb = np.uint8(heatmap_rgb * 255)
+    overlayed = cv2.addWeighted(np.uint8(img * 255), 0.5, heatmap_rgb, 0.5, 0)
+
+    # LIME explanation
+    explanation = explainer.explain_instance(
+        image=img, classifier_fn=predict_fn,
+        top_labels=1, hide_color=0, num_samples=1000
+    )
+    temp, mask = explanation.get_image_and_mask(label=pred_idx, positive_only=True, num_features=10, hide_rest=False)
+
+    # Plot side by side
+    fig, axs = plt.subplots(1, 3, figsize=(15, 5))
+    axs[0].imshow(img)
+    axs[0].set_title(f"Original\nTrue: {true_label}", fontsize=11)
+    axs[1].imshow(overlayed)
+    axs[1].set_title(f"Grad-CAM\nPred: {pred_label}", fontsize=11)
+    axs[2].imshow(mark_boundaries(temp, mask))
+    axs[2].set_title(f"LIME\nPred: {pred_label}", fontsize=11)
+    for ax in axs:
+        ax.axis('off')
+    summary = explanation_text.get(pred_label, "Model detected features matching this class.")
+    plt.figtext(0.5, 0.01, summary, wrap=True, ha='center', fontsize=10)
+    plt.tight_layout(rect=[0, 0.03, 1, 1])
+    st.pyplot(fig)
+    plt.close()
+
+# Streamlit app UI
+st.set_page_config(page_title="🧠 Retina Disease Classifier with Grad-CAM & LIME", layout="centered")
+st.title("🧠 Retina Disease Classifier with Grad-CAM & LIME")
 
 model = load_model()
+last_conv_layer_name = find_last_conv_layer(model)
 
-uploaded_file = st.file_uploader("📤 Upload a retinal image", type=["jpg", "jpeg", "png"])
+uploaded_file = st.file_uploader("Upload a retinal image", type=["jpg", "jpeg", "png"])
 if uploaded_file:
     file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
     bgr_img = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
     rgb_img = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2RGB)
 
-    circ, clahe, sharp, final = preprocess_image(rgb_img)
-
-    st.subheader("🧪 Preprocessing Pipeline (Left ➝ Right)")
-    steps = [
-        ("📷 Original", rgb_img),
-        ("🔵 Circular Crop", circ),
-        ("⚪ CLAHE", clahe),
-        ("🟣 Sharpened", sharp),
-        ("📐 Resized", (final * 255).astype(np.uint8))
-    ]
-    cols = st.columns(len(steps))
-    for col, (label, img) in zip(cols, steps):
-        col.image(img, caption=label, use_container_width=True)
-
-    input_tensor = np.expand_dims(final, axis=0)
+    # Preprocess image
+    processed_img = preprocess_image(rgb_img)
+
+    # Predict
+    input_tensor = np.expand_dims(processed_img, axis=0)
     preds = model.predict(input_tensor)
     if isinstance(preds, dict):
         preds = list(preds.values())[0]
-
     pred_idx = np.argmax(preds)
     pred_label = CLASS_NAMES[pred_idx]
     confidence = np.max(preds) * 100
 
-    st.success(f"✅ Prediction: **{pred_label}**")
-    st.info(f"🔍 Confidence: {confidence:.2f}%")
+    st.success(f"Prediction: **{pred_label}** with confidence {confidence:.2f}%")
 
-    show_gradcam(model, final, pred_idx)
-    show_lime(model, final, pred_idx)
+    # Show Grad-CAM and LIME visualizations
+    display_combined_visualization(processed_img, "Unknown (Uploaded)", pred_label, pred_idx, last_conv_layer_name)