Update app.py

app.py

@@ -1,11 +1,10 @@
 import gradio as gr
 import numpy as np
-from PIL import Image
-import tensorflow as tf
 from tensorflow.keras.models import load_model
 from tensorflow.keras import backend as K
+from PIL import Image
 
-# ---- Custom metrics/loss required to load your model ---- #
+# Custom metric
 def dice_coef(y_true, y_pred):
     y_true = K.cast(y_true, 'float32')
     y_pred = K.cast(y_pred, 'float32')
@@ -14,38 +13,31 @@ def dice_coef(y_true, y_pred):
     intersect = K.sum(y_true_f * y_pred_f)
     return (2. * intersect + K.epsilon()) / (K.sum(y_true_f) + K.sum(y_pred_f) + K.epsilon())
 
-# ---- Load model from Hugging Face model repo ---- #
-model_path = "CNNSegmentation_model.keras"   # <-- Replace with your HF model ID
-model = load_model(model_path, custom_objects={'dice_coef': dice_coef}, compile=False)
+# Global variable for model
+model = None
 
-IMG_SIZE = (224, 224)
+def get_model():
+    global model
+    if model is None:
+        # Load the .keras model only once when first needed
+        model = load_model("CNNSegmentation_model.keras", custom_objects={'dice_coef': dice_coef}, compile=False)
+    return model
 
-# ---- Prediction function for Gradio ---- #
 def predict(img):
-    # Convert to grayscale if not already
-    img = img.convert("L")
-
-    # Resize
-    img_resized = img.resize(IMG_SIZE)
-    x = np.array(img_resized) / 255.0
-    x = np.expand_dims(x, axis=(0, -1))  # shape: (1, 224, 224, 1)
-
-    # Predict mask
+    model = get_model()
+    img = img.convert("L").resize((224, 224))
+    x = np.expand_dims(np.expand_dims(np.array(img)/255.0, -1), 0)
     pred = model.predict(x)[0]
-    mask = (pred > 0.5).astype(np.uint8) * 255  # Convert to binary mask
-
-    # Convert numpy mask to image
-    mask_img = Image.fromarray(mask.squeeze(), mode="L")
-    return mask_img
+    mask = (pred > 0.5).astype(np.uint8) * 255
+    return Image.fromarray(mask.squeeze(), "L")
 
-# ---- Gradio Interface ---- #
 interface = gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="pil", label="Upload Brain MRI Image"),
-    outputs=gr.Image(type="pil", label="Predicted Tumor Mask"),
-    title="Brain Tumor Segmentation",
-    description="Upload a brain MRI image to generate a tumor segmentation mask."
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil"),
+    title="Brain Tumor Segmentation"
 )
 
-# interface.launch()
-interface.launch(server_preprocess=False, ssr_mode=False)
+# Plain launch
+interface.launch()
+