Update app.py (#6)

- Update app.py (fddd2289eb49d81e570be22d6b265f316194ce72)


Co-authored-by: MUHAMMAD IBRAHIM <ibrahim313@users.noreply.huggingface.co>

app.py

@@ -3,117 +3,93 @@ import os, glob, traceback
 import numpy as np
 from PIL import Image
 import gradio as gr
-import tensorflow as tf
 
-# 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
+# Use Keras 3 with the TensorFlow backend
+os.environ.setdefault("KERAS_BACKEND", "tensorflow")
+import keras  # Keras 3.x
+from huggingface_hub import hf_hub_download
 
 HF_MODEL_ID = "Vedag812/xray_cnn"
 CLASS_NAMES = ["NORMAL", "PNEUMONIA"]
 
-
+@gr.cache_resource
 def load_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)
+    # safe_mode=False allows loading models saved with older or custom configs
+    model = keras.saving.load_model(model_path, compile=False, safe_mode=False)
+    return model
 
 def _infer_input_shape(model):
-    """returns (H, W, C) with integers if available, else defaults to (150,150,1)"""
-    shape = None
+    # returns (H, W, C)
     try:
-        shape = tuple(model.inputs[0].shape.as_list())  # works on many TF models
+        shp = tuple(model.inputs[0].shape)
     except Exception:
-        try:
-            shape = tuple(model.input_shape)
-        except Exception:
-            pass
-    if not shape or len(shape) < 4:
+        shp = getattr(model, "input_shape", None)
+        if shp is None:
+            return 150, 150, 1
+    if len(shp) < 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
+    H = int(shp[1]) if shp[1] is not None else 150
+    W = int(shp[2]) if shp[2] is not None else 150
+    C = int(shp[3]) if shp[3] is not None else 1
     return H, W, C
 
-def preprocess(pil_img: Image.Image, target_hw_c):
-    H, W, C = target_hw_c
-    # always start from grayscale so intensity stays consistent
+def preprocess(pil_img: Image.Image, target):
+    H, W, C = target
     g = pil_img.convert("L").resize((W, H))
-    g_arr = np.array(g).astype("float32") / 255.0  # (H,W)
+    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)
+        x = np.expand_dims(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)
+        x = np.expand_dims(np.stack([arr]*3, axis=-1), 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)
+        x = np.expand_dims(np.repeat(arr[..., None], C, axis=-1), 0)
     return x
 
 def predict_fn(pil_img: Image.Image):
     try:
+        if pil_img is None:
+            return {"NORMAL": 0.0, "PNEUMONIA": 0.0}, "Please upload an image or pick a sample."
         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
+        y = model.predict(x, verbose=0)
+        prob = float(np.ravel(y)[0])  # sigmoid output
+        idx = int(prob > 0.5)
+        conf = prob if 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}%"
+        msg = f"Prediction: {CLASS_NAMES[idx]} | Confidence: {conf*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."
+            "If this persists, make sure the Space has keras>=3 and tensorflow>=2.16."
         )
-        err_text = "⚠️ Error during prediction:\n\n" + str(e) + "\n\n" + tip
-        return {"NORMAL": 0.0, "PNEUMONIA": 0.0}, err_text
+        err = f"⚠️ Error during prediction:\n\n{e}\n\n{tip}"
+        # Optional: uncomment next line to print full stack to the Space logs
+        # print(traceback.format_exc())
+        return {"NORMAL": 0.0, "PNEUMONIA": 0.0}, err
 
 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]
+    for pat in ("images/*.jpeg", "images/*.jpg", "images/*.png"):
+        files.extend(glob.glob(pat))
+    return [[p] for p in sorted(files)]
 
 with gr.Blocks(css="""
 .gradio-container {max-width: 980px !important; margin: auto;}
 #title {text-align:center;}
-.card {border:1px solid #e5e7eb; border-radius:16px; padding:16px;}
 """) as demo:
     gr.Markdown("<h1 id='title'>Chest X-Ray Classification</h1>")
-    gr.Markdown("Upload an image or click a sample from the gallery. The model predicts NORMAL or PNEUMONIA.")
+    gr.Markdown("Upload an image or click a sample. The model predicts NORMAL or PNEUMONIA.")
 
     with gr.Row():
         with gr.Column(scale=2):
             inp = gr.Image(type="pil", image_mode="L", label="Upload X-ray")
             with gr.Row():
                 btn = gr.Button("Predict", variant="primary")
-                gr.ClearButton(components=[inp], value="Clear")
+                gr.ClearButton(components=[inp])
             gr.Markdown("### Samples")
-            gr.Examples(
-                examples=list_examples(),
-                inputs=inp,
-                examples_per_page=12,
-            )
+            gr.Examples(examples=list_examples(), inputs=inp, examples_per_page=12)
         with gr.Column(scale=1):
             probs = gr.Label(num_top_classes=2, label="Class probabilities")
             out_text = gr.Markdown()