Update app.py

app.py

@@ -7,7 +7,6 @@ import matplotlib.pyplot as plt
 import matplotlib.cm as cm
 from lime import lime_image
 from skimage.segmentation import mark_boundaries
-from keras.models import load_model as keras_load_model
 from keras.layers import BatchNormalization, DepthwiseConv2D, TFSMLayer
 
 # --- Fix deserialization issues ---
@@ -28,25 +27,18 @@ DepthwiseConv2D.from_config = classmethod(patched_dwconv_from_config)
 # --- Constants ---
 IMG_SIZE = (224, 224)
 CLASS_NAMES = ['Normal', 'Diabetes', 'Glaucoma', 'Cataract', 'AMD', 'Hypertension', 'Myopia', 'Others']
-gradcam_enabled = False  # Default to False
 
-# --- Load model auto-detect ---
+# --- Load model using TFSMLayer ---
 @st.cache_resource
-def load_model_auto(path="Model"):
-    global gradcam_enabled
-    if not os.path.exists(path):
-        st.error(f"❌ Model path '{path}' not found!")
+def load_model():
+    model_path = "Model"  # Folder path to TF SavedModel
+    if not os.path.exists(model_path):
+        st.error(f"❌ Model folder '{model_path}' not found!")
         st.stop()
-
     try:
-        if os.path.isdir(path):
-            model = tf.keras.Sequential([TFSMLayer(path, call_endpoint="serving_default")])
-            gradcam_enabled = False
-        elif path.endswith(('.keras', '.h5')):
-            model = keras_load_model(path)
-            gradcam_enabled = True
-        else:
-            raise ValueError("Unsupported model format. Use .keras, .h5 or SavedModel folder.")
+        model = tf.keras.Sequential([
+            TFSMLayer(model_path, call_endpoint="serving_default")
+        ])
         return model
     except Exception as e:
         st.error(f"❌ Error loading model: {str(e)}")
@@ -74,15 +66,16 @@ def preprocess_and_show_steps(img):
 
     resized = cv2.resize(sharp, IMG_SIZE) / 255.0
 
+    # Show preprocessing stages
     fig, axs = plt.subplots(1, 4, figsize=(20, 5))
     axs[0].imshow(img)
-    axs[0].set_title("Original")
+    axs[0].set_title("Original Image")
     axs[1].imshow(circ)
-    axs[1].set_title("Circular Crop")
+    axs[1].set_title("After Circular Crop")
     axs[2].imshow(clahe_img)
-    axs[2].set_title("CLAHE")
+    axs[2].set_title("After CLAHE")
     axs[3].imshow(resized)
-    axs[3].set_title("Sharpened + Resize")
+    axs[3].set_title("Sharpen + Resize")
     for ax in axs:
         ax.axis("off")
     st.pyplot(fig)
@@ -90,26 +83,6 @@ def preprocess_and_show_steps(img):
 
     return resized
 
-# --- Grad-CAM ---
-def find_last_conv_layer(model):
-    for layer in reversed(model.layers):
-        if isinstance(layer, tf.keras.layers.Conv2D):
-            return layer.name
-    return None
-
-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_array)
-        loss = predictions[:, class_index]
-    grads = tape.gradient(loss, conv_outputs)
-    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
-    conv_outputs = conv_outputs[0]
-    heatmap = conv_outputs @ pooled_grads[..., tf.newaxis]
-    heatmap = tf.squeeze(heatmap)
-    heatmap = tf.maximum(heatmap, 0) / tf.math.reduce_max(heatmap)
-    return heatmap.numpy()
-
 # --- LIME Explainer ---
 explainer = lime_image.LimeImageExplainer()
 def predict_fn(images):
@@ -131,77 +104,55 @@ explanation_text = {
     'Others': "Non-specific features detected, marked as Others."
 }
 
-# --- Display explanations ---
-def display_explanations(img, input_array, pred_label, pred_idx):
-    overlayed = None
-    if gradcam_enabled:
-        try:
-            conv_layer = find_last_conv_layer(model)
-            if conv_layer:
-                heatmap = generate_gradcam(model, input_array, pred_idx, conv_layer)
-                heatmap = cv2.resize(heatmap, IMG_SIZE)
-                heatmap = cv2.GaussianBlur(np.uint8(255 * heatmap), (7, 7), 0)
-                heatmap_rgb = np.uint8(cm.jet(heatmap / 255.0)[..., :3] * 255)
-                overlayed = cv2.addWeighted(np.uint8(img * 255), 0.5, heatmap_rgb, 0.5, 0)
-        except Exception as e:
-            st.warning(f"⚠️ Grad-CAM failed: {e}")
-    else:
-        st.info("ℹ️ Grad-CAM is disabled for this model.")
-
-    explanation = explainer.explain_instance(img, classifier_fn=predict_fn, top_labels=1, hide_color=0, num_samples=1000)
+# --- Display LIME only (Grad-CAM not possible with TFSMLayer) ---
+def display_lime_visualization(img, true_label, pred_label, pred_idx):
+    st.info("⚠️ Grad-CAM is disabled because the model is loaded as a TFSMLayer (inference-only).")
+
+    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)
 
-    fig, axs = plt.subplots(1, 3 if overlayed is not None else 2, figsize=(15, 5))
+    fig, axs = plt.subplots(1, 2, figsize=(12, 5))
     axs[0].imshow(img)
-    axs[0].set_title("Original")
+    axs[0].set_title(f"Original\nTrue: {true_label}")
     axs[1].imshow(mark_boundaries(temp, mask))
-    axs[1].set_title("LIME Explanation")
-    if overlayed is not None:
-        axs[2].imshow(overlayed)
-        axs[2].set_title("Grad-CAM")
+    axs[1].set_title(f"LIME Explanation\nPred: {pred_label}")
     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)
-    st.pyplot(fig)
-    plt.close()
-
-# --- Probability chart ---
-def plot_probabilities(probs, class_names):
-    fig, ax = plt.subplots(figsize=(8, 4))
-    bars = ax.barh(class_names, probs * 100, color='skyblue')
-    ax.set_xlim(0, 100)
-    ax.set_xlabel("Confidence (%)")
-    ax.set_title("Prediction Probabilities")
-    for bar, prob in zip(bars, probs):
-        ax.text(prob * 100 + 1, bar.get_y() + bar.get_height()/2, f"{prob*100:.1f}%", va='center')
+    plt.tight_layout(rect=[0, 0.03, 1, 1])
     st.pyplot(fig)
     plt.close()
 
 # --- Streamlit UI ---
-st.set_page_config(page_title="🧠 Retina Disease Classifier", layout="centered")
-st.title("🧠 Retina Disease Classifier with Grad-CAM & LIME")
+st.set_page_config(page_title="🧠 Retina Disease Classifier with LIME", layout="centered")
+st.title("🧠 Retina Disease Classifier with LIME Explanation")
 
-# Load model
-model = load_model_auto("Model")  # Folder or .keras/.h5 file
+model = load_model()
 
-# Upload image
-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)
 
     processed_img = preprocess_and_show_steps(rgb_img)
-    input_tensor = np.expand_dims(processed_img, axis=0)
 
-    preds = predict_fn(input_tensor)
+    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}** with confidence {confidence:.2f}%")
-    plot_probabilities(preds[0], CLASS_NAMES)
+    st.success(f"Prediction: **{pred_label}** with confidence {confidence:.2f}%")
 
-    display_explanations(processed_img, input_tensor, pred_label, pred_idx)
+    display_lime_visualization(processed_img, "Uploaded Image", pred_label, pred_idx)