Update app.py

app.py

@@ -1,41 +1,40 @@
 import os
-os.environ["TF_USE_LEGACY_KERAS"] = "1"
-os.environ["KERAS_BACKEND"] = "tensorflow"
-os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
-
-import gdown
-import cv2
 import numpy as np
+import cv2
 import tensorflow as tf
 import streamlit as st
+import gdown
 from PIL import Image
-from keras.models import load_model
+from keras.layers import BatchNormalization as KBatchNormalization
 from keras_cv_attention_models.coatnet import CoAtNet0
 
-# ------------------ Settings ------------------
+# Patch BatchNormalization to fix axis deserialization issue
+class PatchedBatchNormalization(KBatchNormalization):
+    @classmethod
+    def from_config(cls, config):
+        if isinstance(config.get("axis"), list):
+            config["axis"] = config["axis"][0]
+        return super().from_config(config)
+
 IMG_SIZE = (224, 224)
 CLASS_NAMES = ['Normal', 'Diabetes', 'Glaucoma', 'Cataract', 'AMD', 'Hypertension', 'Myopia', 'Others']
 
-# ------------------ Load Model ------------------
-@st.cache_resource
-def load_coatnet_model():
+@st.cache_resource(show_spinner=True)
+def load_model():
     model_path = "model.keras"
-
     if not os.path.exists(model_path):
         st.info("📥 Downloading model from Google Drive...")
         url = "https://drive.google.com/uc?id=1Gm2O77uWSUnajL0iFlFJtVk_UEN_wrTN"
         gdown.download(url, model_path, quiet=False, fuzzy=True)
-
-    if os.path.getsize(model_path) < 1_000_000:
-        raise ValueError("❌ Model file is too small. Download may have failed.")
-
+        if os.path.getsize(model_path) < 1_000_000:
+            raise ValueError("❌ Downloaded model is too small. Download might have failed!")
     try:
-        model = load_model(
+        model = tf.keras.models.load_model(
             model_path,
             compile=False,
             custom_objects={
+                "BatchNormalization": PatchedBatchNormalization,
                 "CoAtNet0": CoAtNet0,
-                "Functional": tf.keras.models.Model,
                 "gelu": tf.keras.activations.gelu
             }
         )
@@ -44,9 +43,6 @@ def load_coatnet_model():
         st.error(f"❌ Failed to load model: {e}")
         raise
 
-model = load_coatnet_model()
-
-# ------------------ Image Preprocessing ------------------
 def crop_circle(img):
     h, w = img.shape[:2]
     center = (w // 2, h // 2)
@@ -55,20 +51,20 @@ def crop_circle(img):
     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)
+        mask = np.stack([mask]*3, axis=-1)
     img[~mask] = 0
     return img
 
 def apply_clahe(img):
     lab = cv2.cvtColor(img, cv2.COLOR_RGB2LAB)
     l, a, b = cv2.split(lab)
-    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
+    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)
+    blur = cv2.GaussianBlur(img, (0,0), sigma)
     return cv2.addWeighted(img, 4, blur, -4, 128)
 
 def resize_normalize(img):
@@ -83,19 +79,19 @@ def preprocess_image(img):
     img = resize_normalize(img)
     return img
 
-# ------------------ Streamlit UI ------------------
+# Streamlit UI
 st.set_page_config(page_title="🧠 Retina Disease Classifier", layout="centered")
-
 st.title("🧠 Retina Disease Classifier")
-st.markdown("Upload a retinal image to detect possible diseases using a CoAtNet model.")
+st.markdown("Upload a retinal image and get the predicted disease class using the CoAtNet model.")
+
+model = load_model()
 
-uploaded_file = st.file_uploader("📤 Upload Retinal Image", type=["jpg", "jpeg", "png"])
+uploaded_file = st.file_uploader("📤 Upload Image", type=["jpg", "jpeg", "png"])
 
 if uploaded_file is not None:
     file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
     bgr_img = cv2.imdecode(file_bytes, 1)
     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)
@@ -109,5 +105,5 @@ if uploaded_file is not None:
     st.success(f"✅ **Prediction:** `{pred_label}`")
     st.info(f"🔍 Confidence: **{confidence:.2f}%**")
 
-    st.subheader("🧪 Preprocessed Input")
-    st.image((preprocessed * 255).astype(np.uint8), caption="Model Input", use_column_width=True)
+    st.subheader("🧪 Preprocessed Input to Model")
+    st.image((preprocessed * 255).astype(np.uint8), caption="Preprocessed Image", use_column_width=True)