Update app.py

app.py

@@ -1,5 +1,4 @@
 import os
-# Set Protocol Buffers implementation to Python (must be before TensorFlow imports)
 os.environ['PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION'] = 'python'
 
 import numpy as np
@@ -10,8 +9,11 @@ from keras_cv_attention_models.coatnet import CoAtNet0
 from keras.layers import BatchNormalization, DepthwiseConv2D, Input, TFSMLayer
 from keras.models import Model
 from keras.saving import register_keras_serializable
+import matplotlib.pyplot as plt
+from lime import lime_image
+from skimage.segmentation import mark_boundaries
 
-# --- Fix BatchNormalization axis bug during deserialization ---
+# --- Fix deserialization issues ---
 original_bn_from_config = BatchNormalization.from_config
 def patched_bn_from_config(cls, config, *args, **kwargs):
     if "axis" in config and isinstance(config["axis"], (list, tuple)):
@@ -19,7 +21,6 @@ def patched_bn_from_config(cls, config, *args, **kwargs):
     return original_bn_from_config(config, *args, **kwargs)
 BatchNormalization.from_config = classmethod(patched_bn_from_config)
 
-# --- Fix DepthwiseConv2D deserialization by removing unsupported 'groups' kwarg ---
 original_dwconv_from_config = DepthwiseConv2D.from_config
 def patched_dwconv_from_config(cls, config, *args, **kwargs):
     if "groups" in config:
@@ -27,12 +28,10 @@ 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 Functional Model for deserialization ---
 @register_keras_serializable(package='Custom', name='Functional')
 class Functional(tf.keras.models.Model):
     pass
 
-# --- Register a minimal stub for TFOpLambda layer ---
 @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):
@@ -50,20 +49,15 @@ CLASS_NAMES = ['Normal', 'Diabetes', 'Glaucoma', 'Cataract', 'AMD', 'Hypertensio
 
 @st.cache_resource
 def load_model():
-    model_path = "Model"  # SavedModel directory
-
+    model_path = "Model"
     if not os.path.exists(model_path):
-        st.error(f"❌ Model directory '{model_path}' not found! Please ensure the folder is in the app directory.")
+        st.error(f"❌ Model directory '{model_path}' not found!")
         st.stop()
-
-    st.info("📥 Loading model from SavedModel directory using TFSMLayer...")
-
     try:
         tfsm_layer = TFSMLayer(model_path, call_endpoint="serving_default")
         inputs = Input(shape=(224, 224, 3))
         outputs = tfsm_layer(inputs)
         model = Model(inputs=inputs, outputs=outputs)
-        st.success("✅ Model loaded successfully!")
         return model
     except Exception as e:
         st.error(f"❌ Error loading model: {str(e)}")
@@ -76,11 +70,7 @@ def crop_circle(img):
     Y, X = np.ogrid[:h, :w]
     dist = np.sqrt((X - center[0])**2 + (Y - center[1])**2)
     mask = dist <= radius
-    if img.ndim == 3:
-        mask = np.stack([mask]*3, axis=-1)
-    img = img.copy()
-    img[~mask] = 0
-    return img
+    return cv2.bitwise_and(img, img, mask=mask.astype(np.uint8))
 
 def apply_clahe(img):
     lab = cv2.cvtColor(img, cv2.COLOR_RGB2LAB)
@@ -92,51 +82,79 @@ def apply_clahe(img):
 
 def sharpen_image(img, sigma=10):
     blur = cv2.GaussianBlur(img, (0,0), sigma)
-    sharpened = cv2.addWeighted(img, 4, blur, -4, 128)
-    return sharpened
+    return cv2.addWeighted(img, 4, blur, -4, 128)
 
 def resize_normalize(img):
     img = cv2.resize(img, IMG_SIZE)
-    img = img / 255.0
-    return img
+    return img / 255.0
 
 def preprocess_image(img):
-    img = crop_circle(img)
-    img = apply_clahe(img)
-    img = sharpen_image(img)
-    img = resize_normalize(img)
-    return img
+    circ = crop_circle(img)
+    clahe = apply_clahe(circ)
+    sharp = sharpen_image(clahe)
+    resized = resize_normalize(sharp)
+    return circ, clahe, sharp, resized
+
+def show_step(title, img):
+    st.subheader(title)
+    st.image(img, use_column_width=True)
+
+def show_gradcam(model, img, class_idx):
+    grad_model = Model(model.inputs, [model.layers[-1].output, model.layers[-2].output])
+    img_tensor = tf.convert_to_tensor(img[np.newaxis, ...])
+    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)
+    cam = tf.nn.relu(cam)
+    cam = cam.numpy()
+    cam = cv2.resize(cam, IMG_SIZE)
+    cam = (cam - cam.min()) / (cam.max() - cam.min())
+    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_column_width=True)
+
+def show_lime(model, img, class_idx):
+    explainer = lime_image.LimeImageExplainer()
+    def predict_fn(images):
+        images = np.array(images)
+        return model.predict(images)
+    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_column_width=True)
 
 # --- Streamlit UI ---
 st.set_page_config(page_title="🧠 Retina Disease Classifier", layout="centered")
 st.title("🧠 Retina Disease Classifier")
-st.markdown("Upload a retinal image and get the predicted disease class using a CoAtNet model.")
 
 model = load_model()
-
 uploaded_file = st.file_uploader("📤 Upload a retinal image", type=["jpg", "jpeg", "png"])
 
-if uploaded_file is not None:
+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)
-    st.image(rgb_img, caption="Original Image", use_column_width=True)
 
-    preprocessed = preprocess_image(rgb_img)
-    input_tensor = np.expand_dims(preprocessed, axis=0)
+    show_step("📷 Original Image", rgb_img)
+    circ, clahe, sharp, final = preprocess_image(rgb_img)
+    show_step("🔵 Circular Cropped", circ)
+    show_step("⚪ CLAHE Applied", clahe)
+    show_step("🟣 Sharpened", sharp)
+    show_step("📐 Final Resized", (final * 255).astype(np.uint8))
 
+    input_tensor = np.expand_dims(final, axis=0)
     preds = model.predict(input_tensor)
-
-    # Handle dict output from TFSMLayer
-    if isinstance(preds, dict):
-        preds = list(preds.values())[0]
-
+    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.success(f"✅ Prediction: **{pred_label}**")
     st.info(f"🔍 Confidence: {confidence:.2f}%")
 
-    st.subheader("🧪 Preprocessed Input to Model")
-    st.image((preprocessed * 255).astype(np.uint8), caption="Preprocessed Image", use_column_width=True)
+    show_gradcam(model, final, pred_idx)
+    show_lime(model, final, pred_idx)