Update app.py

app.py

@@ -8,7 +8,6 @@ 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
 
@@ -41,28 +40,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 model using TFSMLayer
+# --- Load SavedModel as TFSMLayer wrapped model ---
 @st.cache_resource
 def load_model():
-    model_path = "Model"  # Adjust path to your SavedModel directory
+    model_path = "Model"  # Your SavedModel directory path
     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=(IMG_SIZE[0], IMG_SIZE[1], 3))
+        inputs = Input(shape=(224, 224, 3))
         outputs = tfsm_layer(inputs)
         model = Model(inputs=inputs, outputs=outputs)
         return model
     except Exception as e:
-        st.error(f"❌ Error loading model: {str(e)}")
+        st.error(f"❌ Error loading model with TFSMLayer: {str(e)}")
         st.stop()
 
-# Preprocessing functions
+# --- Preprocessing functions ---
 def crop_circle(img):
     h, w = img.shape[:2]
     center = (w // 2, h // 2)
@@ -93,113 +92,93 @@ def preprocess_image(img):
     clahe = apply_clahe(circ)
     sharp = sharpen_image(clahe)
     resized = resize_normalize(sharp)
-    return resized
+    return circ, clahe, sharp, resized
 
-# --- Find Last Conv Layer for Grad-CAM ---
-def find_last_conv_layer(model):
+# --- Grad-CAM ---
+def show_gradcam(model, img, class_idx):
+    last_conv_layer = None
     for layer in reversed(model.layers):
-        # 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])
+        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, ...])
+    
     with tf.GradientTape() as tape:
-        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")
+        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")
 
 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)
 
-    # Preprocess image
-    processed_img = preprocess_image(rgb_img)
-
-    # Predict
-    input_tensor = np.expand_dims(processed_img, axis=0)
+    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)
     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}** with confidence {confidence:.2f}%")
+    st.success(f"✅ Prediction: **{pred_label}**")
+    st.info(f"🔍 Confidence: {confidence:.2f}%")
 
-    # Show Grad-CAM and LIME visualizations
-    display_combined_visualization(processed_img, "Unknown (Uploaded)", pred_label, pred_idx, last_conv_layer_name)
+    show_gradcam(model, final, pred_idx)
+    show_lime(model, final, pred_idx)