Update app.py (#3)

- Update app.py (fe73a0757085ed388a81afce43206377986ef5c5)


Co-authored-by: MUHAMMAD IBRAHIM <[email protected]>

app.py

@@ -1,43 +1,98 @@
 # app.py
-import os, glob
+import os, glob, traceback
 import numpy as np
 from PIL import Image
 import gradio as gr
 import tensorflow as tf
-from functools import lru_cache
-from huggingface_hub import hf_hub_download
+
+# try to use Keras 3 loader if present (many models saved with Keras 3 need this)
+KERAS3_AVAILABLE = False
+try:
+    import keras  # pip package "keras" (v3.x)
+    KERAS3_AVAILABLE = int(keras.__version__.split(".")[0]) >= 3
+except Exception:
+    keras = None
 
 HF_MODEL_ID = "Vedag812/xray_cnn"
 CLASS_NAMES = ["NORMAL", "PNEUMONIA"]
 
-@lru_cache(maxsize=1)
+@gr.cache_resource
 def load_model():
-    model_path = hf_hub_download(repo_id=HF_MODEL_ID, filename="xray_cnn.keras")
-    model = tf.keras.models.load_model(model_path, compile=False)
-    return model
+    from huggingface_hub import hf_hub_download
+    model_path = hf_hub_download(HF_MODEL_ID, filename="xray_cnn.keras")
+    # Keras 3 path
+    if KERAS3_AVAILABLE:
+        os.environ.setdefault("KERAS_BACKEND", "tensorflow")
+        try:
+            return keras.saving.load_model(model_path, compile=False, safe_mode=False)
+        except Exception:
+            # fall back to tf.keras if that fails for any reason
+            pass
+    # tf.keras path
+    return tf.keras.models.load_model(model_path, compile=False)
+
+def _infer_input_shape(model):
+    """returns (H, W, C) with integers if available, else defaults to (150,150,1)"""
+    shape = None
+    try:
+        shape = tuple(model.inputs[0].shape.as_list())  # works on many TF models
+    except Exception:
+        try:
+            shape = tuple(model.input_shape)
+        except Exception:
+            pass
+    if not shape or len(shape) < 4:
+        return 150, 150, 1
+    H = int(shape[1]) if shape[1] else 150
+    W = int(shape[2]) if shape[2] else 150
+    C = int(shape[3]) if shape[3] else 1
+    return H, W, C
 
-def preprocess(pil_img: Image.Image):
-    img = pil_img.convert("L").resize((150, 150))
-    arr = np.array(img).astype("float32") / 255.0
-    arr = np.expand_dims(arr, axis=(0, -1))  # shape (1,150,150,1)
-    return arr
+def preprocess(pil_img: Image.Image, target_hw_c):
+    H, W, C = target_hw_c
+    # always start from grayscale so intensity stays consistent
+    g = pil_img.convert("L").resize((W, H))
+    g_arr = np.array(g).astype("float32") / 255.0  # (H,W)
+    if C == 1:
+        x = np.expand_dims(g_arr, axis=(0, -1))       # (1,H,W,1)
+    elif C == 3:
+        x3 = np.stack([g_arr, g_arr, g_arr], axis=-1) # (H,W,3)
+        x = np.expand_dims(x3, axis=0)                # (1,H,W,3)
+    else:
+        # unexpected channel count. tile to that count safely
+        xC = np.repeat(g_arr[..., None], C, axis=-1)
+        x = np.expand_dims(xC, axis=0)
+    return x
 
 def predict_fn(pil_img: Image.Image):
-    model = load_model()
-    x = preprocess(pil_img)
-    prob = float(model.predict(x, verbose=0)[0][0])  # sigmoid
-    pred_idx = int(prob > 0.5)
-    confidence = prob if pred_idx == 1 else 1 - prob
-    probs = {CLASS_NAMES[0]: 1 - prob, CLASS_NAMES[1]: prob}
-    msg = f"Prediction: {CLASS_NAMES[pred_idx]} | Confidence: {confidence*100:.2f}%"
-    return probs, msg
+    try:
+        model = load_model()
+        H, W, C = _infer_input_shape(model)
+        x = preprocess(pil_img, (H, W, C))
+        preds = model.predict(x, verbose=0)
+        # handle models that output shape (1,1) or (1,)
+        prob = float(preds.ravel()[0])
+        pred_idx = int(prob > 0.5)
+        confidence = prob if pred_idx == 1 else 1 - prob
+        probs = {CLASS_NAMES[0]: 1 - prob, CLASS_NAMES[1]: prob}
+        msg = f"Prediction: {CLASS_NAMES[pred_idx]} | Confidence: {confidence*100:.2f}%"
+        return probs, msg
+    except Exception as e:
+        # show a readable error with a tip
+        tip = (
+            "Tip: if this keeps happening, the Space may need keras>=3 to load a model "
+            "saved with newer Keras. I handled both paths here, but if your model was saved "
+            "with a very new version, updating the Space deps can help."
+        )
+        err_text = "⚠️ Error during prediction:\n\n" + str(e) + "\n\n" + tip
+        return {"NORMAL": 0.0, "PNEUMONIA": 0.0}, err_text
 
 def list_examples():
     files = []
     for pattern in ["images/*.jpeg", "images/*.jpg", "images/*.png"]:
         files.extend(glob.glob(pattern))
     files = sorted(files)
-    return [[p] for p in files]  # gr.Examples expects list of [path]
+    return [[p] for p in files]
 
 with gr.Blocks(css="""
 .gradio-container {max-width: 980px !important; margin: auto;}
@@ -52,7 +107,7 @@ with gr.Blocks(css="""
             inp = gr.Image(type="pil", image_mode="L", label="Upload X-ray")
             with gr.Row():
                 btn = gr.Button("Predict", variant="primary")
-                clr = gr.ClearButton(components=[inp], value="Clear")
+                gr.ClearButton(components=[inp], value="Clear")
             gr.Markdown("### Samples")
             gr.Examples(
                 examples=list_examples(),
@@ -63,9 +118,7 @@ with gr.Blocks(css="""
             probs = gr.Label(num_top_classes=2, label="Class probabilities")
             out_text = gr.Markdown()
 
-    # Run on click
     btn.click(predict_fn, inputs=inp, outputs=[probs, out_text])
-    # Also auto-run when image changes (from upload or example click)
     inp.change(predict_fn, inputs=inp, outputs=[probs, out_text])
 
 if __name__ == "__main__":