Add Gradio app and requirements for WBC ResNet50

app.py

@@ -11,22 +11,31 @@ from huggingface_hub import snapshot_download
 MODEL_ID = "varshithkumar/wbc_resnet50"   # your model repo id
 CLASS_NAMES = ['Basophil', 'Eosinophil', 'Lymphocyte', 'Monocyte', 'Neutrophil']
 
+# Globals
+model = None
+infer = None   # serving function
+
 # === Load model ===
 def load_model():
+    global infer
     hf_token = os.environ.get("HF_TOKEN")  # set in Space secrets if repo is private
     try:
         print(f"⏳ Downloading model from Hugging Face Hub: {MODEL_ID}")
         repo_dir = snapshot_download(repo_id=MODEL_ID, repo_type="model", token=hf_token)
         print("✅ Model snapshot downloaded at:", repo_dir)
 
-        # Try loading with tf.saved_model.load instead of keras
+        # Load TF SavedModel
         model = tf.saved_model.load(repo_dir)
         print("✅ Model loaded using tf.saved_model.load()")
-        # 🔍 Debug: print signatures and IO structure
-        print("Available signatures:", list(model.signatures.keys()))
+
+        # Get serving function
         infer = model.signatures["serving_default"]
+
+        # 🔍 Debug info
+        print("Available signatures:", list(model.signatures.keys()))
         print("Serving function inputs:", infer.structured_input_signature)
         print("Serving function outputs:", infer.structured_outputs)
+
         return model
     except Exception as e:
         print("❌ Failed to load model:", e)
@@ -35,7 +44,7 @@ def load_model():
 
 model = load_model()
 if model is None:
-    print("WARNING: Model failed to load. Predictions will return zeros.")
+    print("WARNING: Model failed to load. Predictions will return an error.")
 
 # === preprocessing & prediction ===
 def preprocess_image(img: Image.Image):
@@ -46,30 +55,27 @@ def preprocess_image(img: Image.Image):
     return arr
 
 def predict(image):
-    if model is None:
-        return {cls: 0.0 for cls in CLASS_NAMES}
+    global infer
+    if infer is None:
+        return {"error": "Model not loaded. Check Space logs."}
     try:
         arr = preprocess_image(image)
+        preds = infer(input_layer_2=tf.constant(arr))["output_0"].numpy()
 
-        # SavedModel inference requires calling the serving function
-        infer = model.signatures["serving_default"]
-        preds = infer(tf.constant(arr))  # run inference
-        preds = list(preds.values())[0].numpy()  # extract tensor
         probs = preds[0].tolist()
-
         if len(probs) == len(CLASS_NAMES):
             out = {CLASS_NAMES[i]: float(probs[i]) for i in range(len(CLASS_NAMES))}
         else:
-            out = {f"class_{i}": float(p) for i, p in enumerate(probs)}
+            out = {"class_" + str(i): float(p) for i, p in enumerate(probs)}
         return out
     except Exception as e:
         print("Prediction error:", e)
         traceback.print_exc()
-        return {cls: 0.0 for cls in CLASS_NAMES}
+        return {"error": str(e)}
 
 # === Gradio UI ===
 title = "WBC ResNet50 - White Blood Cell Classifier"
-description = "Upload a blood-smear image. Model resizes input to 224×224. If model fails to load, all predictions will be 0."
+description = "Upload a blood-smear image. Model resizes input to 224×224. If model fails to load, predictions will error."
 
 demo = gr.Interface(
     fn=predict,