Update app.py

app.py

@@ -3,14 +3,11 @@ import numpy as np
 import cv2
 import tensorflow as tf
 import streamlit as st
-from keras.layers import BatchNormalization, DepthwiseConv2D, Input
-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
+from keras.layers import BatchNormalization, DepthwiseConv2D
 
 # --- Fix deserialization issues ---
 original_bn_from_config = BatchNormalization.from_config
@@ -27,75 +24,49 @@ def patched_dwconv_from_config(cls, config, *args, **kwargs):
     return original_dwconv_from_config(config, *args, **kwargs)
 DepthwiseConv2D.from_config = classmethod(patched_dwconv_from_config)
 
-@register_keras_serializable(package='Custom', name='Functional')
-class Functional(tf.keras.models.Model): pass
-
-@register_keras_serializable(package='Custom', name='TFOpLambda')
-class CustomTFOpLambda(tf.keras.layers.Layer):
-    def __init__(self, name=None, trainable=False, dtype=None, function=None, **kwargs):
-        super().__init__(name=name, trainable=trainable, dtype=dtype, **kwargs)
-        self.function = function
-    def call(self, inputs): return inputs
-    def get_config(self):
-        config = super().get_config()
-        config.update({"function": self.function})
-        return config
-
 # --- Constants ---
 IMG_SIZE = (224, 224)
 CLASS_NAMES = ['Normal', 'Diabetes', 'Glaucoma', 'Cataract', 'AMD', 'Hypertension', 'Myopia', 'Others']
 
-# --- Load model with TFSMLayer ---
+# --- Load model function ---
 @st.cache_resource
 def load_model():
-    model_path = "Model"  # Your SavedModel directory path
+    model_path = "Model"  # adjust path to your model folder or file
     if not os.path.exists(model_path):
-        st.error(f"❌ Model directory '{model_path}' not found!")
+        st.error(f"❌ Model directory or file '{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))
-        outputs = tfsm_layer(inputs)
-        model = Model(inputs=inputs, outputs=outputs)
+        model = tf.keras.models.load_model(model_path)
         return model
     except Exception as e:
         st.error(f"❌ Error loading model: {str(e)}")
         st.stop()
 
-# --- Preprocessing functions ---
-def crop_circle(img):
+# --- Preprocessing ---
+def preprocess_image(img):
+    # Crop circular mask
     h, w = img.shape[:2]
-    center = (w // 2, h // 2)
+    center = (w//2, h//2)
     radius = min(center[0], center[1])
     Y, X = np.ogrid[:h, :w]
     dist = np.sqrt((X - center[0])**2 + (Y - center[1])**2)
     mask = dist <= radius
-    return cv2.bitwise_and(img, img, mask=mask.astype(np.uint8))
-
-def apply_clahe(img):
-    lab = cv2.cvtColor(img, cv2.COLOR_RGB2LAB)
+    circ = cv2.bitwise_and(img, img, mask=mask.astype(np.uint8))
+    # CLAHE
+    lab = cv2.cvtColor(circ, cv2.COLOR_RGB2LAB)
     l, a, b = cv2.split(lab)
     clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
     cl = clahe.apply(l)
     merged = cv2.merge((cl,a,b))
-    return cv2.cvtColor(merged, cv2.COLOR_LAB2RGB)
-
-def sharpen_image(img, sigma=10):
-    blur = cv2.GaussianBlur(img, (0,0), sigma)
-    return cv2.addWeighted(img, 4, blur, -4, 128)
-
-def resize_normalize(img):
-    img = cv2.resize(img, IMG_SIZE)
-    return img / 255.0
-
-def preprocess_image(img):
-    circ = crop_circle(img)
-    clahe = apply_clahe(circ)
-    sharp = sharpen_image(clahe)
-    resized = resize_normalize(sharp)
+    clahe_img = cv2.cvtColor(merged, cv2.COLOR_LAB2RGB)
+    # Sharpen
+    blur = cv2.GaussianBlur(clahe_img, (0,0), 10)
+    sharp = cv2.addWeighted(clahe_img, 4, blur, -4, 128)
+    # Resize + normalize
+    resized = cv2.resize(sharp, IMG_SIZE) / 255.0
     return resized
 
-# --- Find last Conv layer for Grad-CAM ---
+# --- Find last Conv2D or MHSA output layer ---
 def find_last_conv_layer(model):
     for layer in reversed(model.layers):
         if isinstance(layer, tf.keras.layers.Conv2D) or 'mhsa_output' in layer.name:
@@ -109,7 +80,7 @@ def generate_gradcam(model, img_array, class_index, layer_name):
         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))
+    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)
@@ -124,7 +95,7 @@ def predict_fn(images):
         preds = list(preds.values())[0]
     return preds
 
-# --- Explanation text ---
+# --- Explanation texts ---
 explanation_text = {
     'Normal': "Model predicted Normal based on healthy optic disc and macula.",
     'Diabetes': "Detected retinal blood vessel changes suggestive of Diabetes.",
@@ -136,22 +107,23 @@ explanation_text = {
     'Others': "Non-specific features detected, marked as Others."
 }
 
-# --- Visualization function ---
+# --- Visualization ---
 def display_combined_visualization(img, true_label, pred_label, pred_idx, layer_name):
     input_array = np.expand_dims(img, axis=0)
 
     # Grad-CAM heatmap
-    try:
-        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)
-    except Exception as e:
-        overlayed = None
-        st.warning(f"⚠️ Grad-CAM generation failed: {e}")
+    overlayed = None
+    if layer_name is not None:
+        try:
+            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)
+        except Exception as e:
+            st.warning(f"⚠️ Grad-CAM generation failed: {e}")
 
     # LIME explanation
     explanation = explainer.explain_instance(
@@ -160,7 +132,7 @@ def display_combined_visualization(img, true_label, pred_label, pred_idx, layer_
     )
     temp, mask = explanation.get_image_and_mask(label=pred_idx, positive_only=True, num_features=10, hide_rest=False)
 
-    # Plot side-by-side
+    # Plot
     cols = 3 if overlayed is not None else 2
     fig, axs = plt.subplots(1, cols, figsize=(15, 5))
     axs[0].imshow(img)
@@ -183,16 +155,15 @@ def display_combined_visualization(img, true_label, pred_label, pred_idx, layer_
     st.pyplot(fig)
     plt.close()
 
-# --- Streamlit UI ---
+# --- Streamlit App ---
 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()
 
-# Try find last conv layer, disable Grad-CAM if not found
 try:
     last_conv_layer_name = find_last_conv_layer(model)
-except ValueError:
+except Exception:
     last_conv_layer_name = None
     st.warning("⚠️ No Conv2D layer found; Grad-CAM will be disabled.")
 
@@ -214,4 +185,4 @@ if uploaded_file:
 
     st.success(f"Prediction: **{pred_label}** with confidence {confidence:.2f}%")
 
-    display_combined_visualization(processed_img, "Unknown (Uploaded)", pred_label, pred_idx, last_conv_layer_name)
+    display_combined_visualization(processed_img, "Uploaded Image", pred_label, pred_idx, last_conv_layer_name)