Update app.py

app.py

@@ -1,63 +1,103 @@
 import gradio as gr
 import numpy as np
 import tensorflow as tf
-from tensorflow.keras.models import Sequential, load_model
-from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout
+from tensorflow.keras.models import Model, load_model
+from tensorflow.keras.layers import GlobalAveragePooling2D, Dense, Dropout
+from tensorflow.keras.preprocessing.image import ImageDataGenerator
 import os
 
-MODEL_FILE = "brain_tumor_cnn.h5"
-IMG_SIZE = (128, 128)
+# Settings
+MODEL_FILE = "brain_tumor_mobilenet.h5"
+IMG_SIZE = (224, 224)
+BATCH_SIZE = 16
+EPOCHS = 5  # change/fine-tune more if you have real data
 
-# -------------------------
-# 1. Train Model (if not exists)
-# -------------------------
+# Utility: prepare data
+def get_dataset(data_dir="data"):
+    """
+    Expect data_dir with two subfolders: `tumor` and `no_tumor`, containing images.
+    """
+    datagen = ImageDataGenerator(
+        rescale=1./255,
+        validation_split=0.2
+    )
+    train_gen = datagen.flow_from_directory(
+        data_dir,
+        target_size=IMG_SIZE,
+        batch_size=BATCH_SIZE,
+        class_mode="binary",
+        subset="training"
+    )
+    val_gen = datagen.flow_from_directory(
+        data_dir,
+        target_size=IMG_SIZE,
+        batch_size=BATCH_SIZE,
+        class_mode="binary",
+        subset="validation"
+    )
+    return train_gen, val_gen
+
+# 1. Build or Load model
 if not os.path.exists(MODEL_FILE):
-    # ⚠️ Dummy training with random data (replace with actual dataset for real use)
-    X_train = np.random.rand(100, IMG_SIZE[0], IMG_SIZE[1], 3)
-    y_train = np.random.randint(0, 2, 100)
+    # if you have dataset
+    has_data = os.path.isdir("data")
+    base_model = tf.keras.applications.MobileNetV2(
+        weights="imagenet",
+        include_top=False,
+        input_shape=(IMG_SIZE[0], IMG_SIZE[1], 3)
+    )
+    x = base_model.output
+    x = GlobalAveragePooling2D()(x)
+    x = Dropout(0.5)(x)
+    output = Dense(1, activation="sigmoid")(x)
+    model = Model(inputs=base_model.input, outputs=output)
 
-    model = Sequential([
-        Conv2D(32, (3,3), activation="relu", input_shape=(IMG_SIZE[0], IMG_SIZE[1], 3)),
-        MaxPooling2D(2,2),
-        Conv2D(64, (3,3), activation="relu"),
-        MaxPooling2D(2,2),
-        Flatten(),
-        Dense(128, activation="relu"),
-        Dropout(0.5),
-        Dense(1, activation="sigmoid")
-    ])
+    # Freeze base layers first
+    for layer in base_model.layers:
+        layer.trainable = False
 
     model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])
-    model.fit(X_train, y_train, epochs=2, batch_size=8, verbose=1)
 
-    model.save(MODEL_FILE)
+    if has_data:
+        train_gen, val_gen = get_dataset("data")
+        model.fit(train_gen, validation_data=val_gen, epochs=EPOCHS)
+        # Optionally unfreeze some layers and fine‐tune
+    else:
+        # Dummy training if no dataset, just random noise (NOT for real use)
+        dummy_x = np.random.rand(20, IMG_SIZE[0], IMG_SIZE[1], 3)
+        dummy_y = np.random.randint(0, 2, size=(20,))
+        model.fit(dummy_x, dummy_y, epochs=2, batch_size=4)
 
-# -------------------------
-# 2. Load model
-# -------------------------
-model = load_model(MODEL_FILE)
+    model.save(MODEL_FILE)
+else:
+    model = load_model(MODEL_FILE)
 
-# -------------------------
-# 3. Prediction Function
-# -------------------------
+# 2. Prediction function
 def predict_brain_tumor(image):
-    image = image.resize(IMG_SIZE)
-    img_array = np.array(image) / 255.0
-    img_array = np.expand_dims(img_array, axis=0)  # (1, 128, 128, 3)
-
-    prediction = model.predict(img_array)[0][0]
-
-    if prediction > 0.5:
-        return "🧠 Tumor Detected"
+    """
+    Input: PIL image via Gradio upload
+    Output: prediction string + probability
+    """
+    if image is None:
+        return "No image provided"
+    img = image.resize(IMG_SIZE)
+    img_arr = np.array(img) / 255.0
+    if img_arr.shape[-1] == 1:
+        img_arr = np.stack([img_arr]*3, axis=-1)  # grayscale → 3 channels
+    elif img_arr.shape[-1] == 4:
+        # drop alpha
+        img_arr = img_arr[..., :3]
+    img_batch = np.expand_dims(img_arr, axis=0)
+    prob = model.predict(img_batch)[0][0]
+    if prob > 0.5:
+        return f"🧠 Tumor Detected (confidence {prob:.2f})"
     else:
-        return "✅ No Tumor Detected"
+        return f"✅ No Tumor Detected (confidence {1-prob:.2f})"
 
-# -------------------------
-# 4. Gradio UI
-# -------------------------
+# 3. Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# 🧠 Brain Tumor Detection (CNN on MRI Images)")
-    gr.Markdown("Upload an MRI image to check if a brain tumor is detected.")
+    gr.Markdown("# Brain Tumor Detection via Transfer Learning (MobileNetV2)")
+    gr.Markdown("Upload an MRI brain scan to detect presence of tumor vs no tumor.")
 
     image_input = gr.Image(type="pil", label="Upload MRI Image")
     output_text = gr.Textbox(label="Prediction")