fixed the layers

app.py

@@ -3,24 +3,26 @@ import tensorflow as tf
 import numpy as np
 from PIL import Image
 
-# --- 1. BUILD X-RAY MODEL (DenseNet121 in 3 Sequential Layers) ---
+# --- 1. BUILD X-RAY MODEL (Matching the 3-layer .h5 file) ---
 def build_xray_model():
+    # include_top=False + pooling=None = 1st layer
     base_model = tf.keras.applications.DenseNet121(
         input_shape=(320, 320, 3), 
         include_top=False, 
-        weights=None
+        weights=None,
+        pooling=None 
     )
     
-    # Wrap in Sequential to match the "3 saved layers" format exactly
+    # We build a Sequential model with exactly 3 layers to match your weights file
     model = tf.keras.Sequential([
-        base_model,
-        tf.keras.layers.GlobalAveragePooling2D(),
-        tf.keras.layers.Dense(14, activation='sigmoid')
+        base_model,                          # Layer 1
+        tf.keras.layers.GlobalAveragePooling2D(), # Layer 2
+        tf.keras.layers.Dense(14, activation='sigmoid') # Layer 3
     ])
     
     try:
         model.load_weights("xray.h5")
-        print("X-Ray weights loaded successfully into Sequential DenseNet121!")
+        print("X-Ray weights loaded successfully (3/3 layers matched)!")
         return model
     except Exception as e:
         print(f"Error loading X-Ray weights: {e}")
@@ -28,6 +30,7 @@ def build_xray_model():
 
 # --- 2. LOAD MODELS ---
 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:
@@ -48,30 +51,36 @@ xray_labels = [
 def predict(img, model_type):
     if img is None: return {"No image": 0.0}
     
-    if model_type == "MRI":
-        if mri_model is None: return {"MRI Model Error": 0.0}
-        
-        # MRI PREPROCESSING: Grayscale (1 channel) and 256x256
-        img = img.convert("L").resize((256, 256))
-        img_array = np.array(img).astype('float32')
-        img_array = np.expand_dims(img_array, axis=(0, -1)) # Shape becomes (1, 256, 256, 1)
-        model, labels = mri_model, mri_labels
-        
-    else:
-        if xray_model is None: return {"X-Ray Model Error": 0.0}
-        
-        # X-RAY PREPROCESSING: RGB (3 channels) and 320x320
-        img = img.convert("RGB").resize((320, 320))
-        img_array = np.array(img).astype('float32')
-        img_array = np.expand_dims(img_array, axis=0) # Shape is (1, 320, 320, 3)
-        model, labels = xray_model, xray_labels
+    try:
+        if model_type == "MRI":
+            if mri_model is None: return {"MRI Model Error": 0.0}
+            
+            # STrictly convert to Grayscale (L) and resize
+            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))
+            model, labels = mri_model, mri_labels
+            
+        else:
+            if xray_model is None: return {"X-Ray Model Error": 0.0}
+            
+            # X-Ray expects RGB (3 channels) and 320x320
+            img = img.convert("RGB").resize((320, 320))
+            img_array = np.array(img).astype('float32')
+            img_array = np.expand_dims(img_array, axis=0)
+            model, labels = xray_model, xray_labels
 
-    # Normalize pixel values
-    img_array /= 255.0
+        # Normalize 0-255 to 0-1
+        img_array /= 255.0
+        
+        # Predict
+        preds = model.predict(img_array)[0]
+        return {labels[i]: float(preds[i]) for i in range(len(labels))}
     
-    # 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}
 
 # --- 5. GRADIO UI ---
 with gr.Blocks() as demo:
@@ -79,15 +88,16 @@ with gr.Blocks() as demo:
     
     with gr.Tabs():
         with gr.TabItem("Brain MRI (256x256)"):
-            mri_in = gr.Image(type="pil")
-            mri_out = gr.Label(num_top_classes=1)
+            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")
             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")
-            xray_out = gr.Label(num_top_classes=1)
+            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")
             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())