Update app.py

app.py

@@ -10,41 +10,42 @@ from sklearn.metrics import accuracy_score
 import joblib
 import os
 
-# ---------------------------
-# 1. Train or Load Model
-# ---------------------------
 MODEL_FILE = "brain_tumor_model.pkl"
 
 if not os.path.exists(MODEL_FILE):
     np.random.seed(42)
     n_samples = 500
 
-    # Simulated dataset (replace with MRI-extracted features later)
+    # Simulated dataset
     X = pd.DataFrame({
         "tumor_size": np.random.uniform(1, 10, n_samples),      # cm
-        "texture_score": np.random.uniform(0, 1, n_samples),    # extracted from MRI texture
-        "age": np.random.randint(20, 80, n_samples),            # patient age
+        "texture_score": np.random.uniform(0, 1, n_samples),    # texture feature
+        "age": np.random.randint(20, 80, n_samples),            # age
         "contrast_intensity": np.random.uniform(50, 200, n_samples), # MRI contrast
     })
 
-    # Simulated binary labels: 1 = Tumor present, 0 = No tumor
+    # Create balanced tumor labels
     y = (
         (X["tumor_size"] > 5).astype(int) |
         (X["contrast_intensity"] > 120).astype(int)
     )
-    y = np.where(np.random.rand(n_samples) > 0.9, 1 - y, y)  # add noise
+
+    # Force balance: make half tumor, half no tumor
+    half = n_samples // 2
+    y[:half] = 0
+    y[half:] = 1
+
+    print("Class distribution:", np.bincount(y))
 
     X_train, X_test, y_train, y_test = train_test_split(
-        X, y, test_size=0.2, random_state=42
+        X, y, test_size=0.2, random_state=42, stratify=y
     )
 
-    pipeline = Pipeline(
-        [
-            ("scaler", StandardScaler()),
-            ("select", SelectKBest(score_func=f_classif, k=3)),  # select best 3 features
-            ("classifier", RandomForestClassifier(n_estimators=100, random_state=42)),
-        ]
-    )
+    pipeline = Pipeline([
+        ("scaler", StandardScaler()),
+        ("select", SelectKBest(score_func=f_classif, k=3)),
+        ("classifier", RandomForestClassifier(n_estimators=200, random_state=42)),
+    ])
 
     pipeline.fit(X_train, y_train)
 
@@ -53,20 +54,14 @@ if not os.path.exists(MODEL_FILE):
 
     joblib.dump(pipeline, MODEL_FILE)
 
-# Load trained model
+# Load model
 model = joblib.load(MODEL_FILE)
 
-# ---------------------------
-# 2. Prediction Function
-# ---------------------------
 def predict_brain_tumor(tumor_size, texture_score, age, contrast_intensity):
     input_data = np.array([[tumor_size, texture_score, age, contrast_intensity]])
     pred = model.predict(input_data)[0]
     return "🧠 Tumor Detected" if pred == 1 else "✅ No Tumor Detected"
 
-# ---------------------------
-# 3. Gradio UI
-# ---------------------------
 with gr.Blocks() as demo:
     gr.Markdown("# 🧠 Brain Tumor Detection")
     gr.Markdown("Predict brain tumor presence using MRI-based features (simulated demo).")