Modified the wrapper

app.py

@@ -3,34 +3,33 @@ import tensorflow as tf
 import numpy as np
 from PIL import Image
 
-# --- 1. BUILD X-RAY MODEL (Matching the 3-layer .h5 file) ---
+# --- 1. X-RAY MODEL RECONSTRUCTION ---
+# Rebuilding exactly based on EfficientNetB1 and 128x128
 def build_xray_model():
-    # include_top=False + pooling=None = 1st layer
-    base_model = tf.keras.applications.DenseNet121(
-        input_shape=(320, 320, 3), 
+    base_model = tf.keras.applications.EfficientNetB1(
+        input_shape=(128, 128, 3), 
         include_top=False, 
-        weights=None,
-        pooling=None 
+        weights=None
     )
     
-    # We build a Sequential model with exactly 3 layers to match your weights file
+    # Built as a Sequential model to perfectly match the 3 saved layers in your .h5 file
     model = tf.keras.Sequential([
-        base_model,                          # Layer 1
-        tf.keras.layers.GlobalAveragePooling2D(), # Layer 2
-        tf.keras.layers.Dense(14, activation='sigmoid') # Layer 3
+        base_model,                                     # Layer 1: Backbone
+        tf.keras.layers.GlobalAveragePooling2D(),       # Layer 2: Pooling
+        tf.keras.layers.Dense(14, activation='sigmoid') # Layer 3: Output head
     ])
     
     try:
         model.load_weights("xray.h5")
-        print("X-Ray weights loaded successfully (3/3 layers matched)!")
+        print("X-Ray weights (EfficientNetB1) loaded successfully!")
         return model
     except Exception as e:
         print(f"Error loading X-Ray weights: {e}")
         return None
 
-# --- 2. LOAD MODELS ---
+# --- 2. LOAD MRI MODEL ---
+# The zaahaa notebook outputs a standard .keras file, so we load it whole
 try:
-    # compile=False helps avoid versioning issues with the optimizer state
     mri_model = tf.keras.models.load_model("mri.keras", compile=False)
     print("MRI model loaded successfully!")
 except Exception as e:
@@ -39,7 +38,7 @@ except Exception as e:
 
 xray_model = build_xray_model()
 
-# --- 3. CONFIGURATION ---
+# --- 3. LABELS ---
 mri_labels = ['Glioma', 'Meningioma', 'Pituitary tumor', 'no tumor']
 xray_labels = [
     'Cardiomegaly', 'Emphysema', 'Effusion', 'Hernia', 'Infiltration', 
@@ -49,55 +48,52 @@ xray_labels = [
 
 # --- 4. PREDICTION LOGIC ---
 def predict(img, model_type):
-    if img is None: return {"No image": 0.0}
+    if img is None: return {"No image provided": 0.0}
     
     try:
         if model_type == "MRI":
-            if mri_model is None: return {"MRI Model Error": 0.0}
+            if mri_model is None: return {"MRI Model Error - Check Logs": 0.0}
             
-            # STrictly convert to Grayscale (L) and resize
+            # The zaahaa notebook uses Grayscale images. We force Grayscale (L) and 256x256.
             img = img.convert("L").resize((256, 256))
             img_array = np.array(img).astype('float32')
-            
-            # Reshape to (1, 256, 256, 1) to match "expected 1 channel"
-            img_array = img_array.reshape((1, 256, 256, 1))
+            img_array = img_array.reshape((1, 256, 256, 1)) # Explicitly format to 1 channel
             model, labels = mri_model, mri_labels
             
         else:
-            if xray_model is None: return {"X-Ray Model Error": 0.0}
+            if xray_model is None: return {"X-Ray Model Error - Check Logs": 0.0}
             
-            # X-Ray expects RGB (3 channels) and 320x320
-            img = img.convert("RGB").resize((320, 320))
+            # Your EfficientNetB1 code used RGB images at 128x128
+            img = img.convert("RGB").resize((128, 128))
             img_array = np.array(img).astype('float32')
-            img_array = np.expand_dims(img_array, axis=0)
+            img_array = np.expand_dims(img_array, axis=0) # Format to 3 channels
             model, labels = xray_model, xray_labels
 
-        # Normalize 0-255 to 0-1
+        # Standard normalization used in both Kaggle notebooks
         img_array /= 255.0
         
-        # Predict
         preds = model.predict(img_array)[0]
         return {labels[i]: float(preds[i]) for i in range(len(labels))}
     
     except Exception as e:
-        return {f"Error during prediction: {str(e)}": 0.0}
+        return {f"Prediction Error: {str(e)}": 0.0}
 
-# --- 5. GRADIO UI ---
+# --- 5. UI APP ---
 with gr.Blocks() as demo:
     gr.Markdown("# 🏥 BTech Medical Diagnostic API")
+    gr.Markdown("Upload an image to get a diagnostic prediction.")
     
     with gr.Tabs():
-        with gr.TabItem("Brain MRI (256x256)"):
-            mri_in = gr.Image(type="pil", label="Upload MRI")
-            mri_out = gr.Label(num_top_classes=1, label="Detection Result")
-            mri_btn = gr.Button("Analyze MRI")
+        with gr.TabItem("Brain MRI Classifier"):
+            mri_in = gr.Image(type="pil", label="Upload Brain MRI")
+            mri_out = gr.Label(num_top_classes=1, label="Result")
+            mri_btn = gr.Button("Analyze MRI", variant="primary")
             mri_btn.click(fn=lambda i: predict(i, "MRI"), inputs=mri_in, outputs=mri_out, api_name="predict_mri")
             
-        with gr.TabItem("Chest X-Ray (320x320)"):
-            xray_in = gr.Image(type="pil", label="Upload X-Ray")
-            xray_out = gr.Label(num_top_classes=1, label="Detection Result")
-            xray_btn = gr.Button("Analyze X-Ray")
+        with gr.TabItem("Chest X-Ray Classifier"):
+            xray_in = gr.Image(type="pil", label="Upload Chest X-Ray")
+            xray_out = gr.Label(num_top_classes=1, label="Result")
+            xray_btn = gr.Button("Analyze X-Ray", variant="primary")
             xray_btn.click(fn=lambda i: predict(i, "X-Ray"), inputs=xray_in, outputs=xray_out, api_name="predict_xray")
 
-# Launch with theme
 demo.launch(theme=gr.themes.Soft())