Update app.py

app.py

@@ -1,77 +1,81 @@
-import streamlit as st
-import torch
-import torch.nn as nn
-from torchvision import transforms
-from PIL import Image
-from TumorModel import TumorClassification, GliomaStageModel
-
-# Title
-st.set_page_config(page_title="Brain Tumor Predictor", layout="centered")
-st.title("🧠 Brain Tumor Classifier & Glioma Stage Predictor")
-
-# Load models
-@st.cache_resource
-def load_models():
-    tumor_model = TumorClassification()
-    tumor_model.load_state_dict(torch.load("BTD_model.pth", map_location=torch.device("cpu")))
-    tumor_model.eval()
-
-    glioma_model = GliomaStageModel()
-    glioma_model.load_state_dict(torch.load("glioma_stages.pth", map_location=torch.device("cpu")))
-    glioma_model.eval()
-
-    return tumor_model, glioma_model
-
-tumor_model, glioma_model = load_models()
-
-# Labels
-tumor_labels = ['glioma', 'meningioma', 'notumor', 'pituitary']
-stage_labels = ['Stage 1', 'Stage 2', 'Stage 3', 'Stage 4']
-
-# Image Preprocessing
-transform = transforms.Compose([
-    transforms.Grayscale(),
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize([0.5], [0.5])
-])
-
-# Tabs
-tab1, tab2 = st.tabs(["🧠 Tumor Type Detection", "🧬 Glioma Stage Prediction"])
-
-# Tab 1: Tumor Classification
-with tab1:
-    st.header("Upload Brain MRI")
-    image_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
-    if image_file:
-        image = Image.open(image_file).convert("L")
-        st.image(image, caption="Uploaded Image", use_column_width=True)
-
-        image_tensor = transform(image).unsqueeze(0)
-        with torch.no_grad():
-            prediction = tumor_model(image_tensor)
-            pred_label = tumor_labels[torch.argmax(prediction)]
-        st.success(f"🧠 Tumor Prediction: **{pred_label.upper()}**")
-
-# Tab 2: Glioma Stage Prediction
-with tab2:
-    st.header("Patient & Genetic Info")
-    gender = st.radio("Gender", ["Male", "Female"])
-    age = st.slider("Age", 1, 100, 25)
-
-    st.subheader("Gene Mutations (0 = No, 1 = Yes)")
-    idh1 = st.radio("IDH1", [0, 1], horizontal=True)
-    tp53 = st.radio("TP53", [0, 1], horizontal=True)
-    atrx = st.radio("ATRX", [0, 1], horizontal=True)
-    pten = st.radio("PTEN", [0, 1], horizontal=True)
-    egfr = st.radio("EGFR", [0, 1], horizontal=True)
-    cic = st.radio("CIC", [0, 1], horizontal=True)
-    pik3ca = st.radio("PIK3CA", [0, 1], horizontal=True)
-
-    if st.button("🔍 Predict Glioma Stage"):
-        gender_val = 0 if gender == "Male" else 1
-        features = torch.tensor([[gender_val, age, idh1, tp53, atrx, pten, egfr, cic, pik3ca]], dtype=torch.float32)
-        with torch.no_grad():
-            prediction = glioma_model(features)
-            stage = stage_labels[torch.argmax(prediction)]
-        st.success(f"🧬 Predicted Glioma Stage: **{stage}**")
+import os
+
+# ✅ Fix permission issues on Hugging Face Spaces
+os.environ["STREAMLIT_HOME"] = "/home/user/app"
+os.environ["XDG_CONFIG_HOME"] = "/home/user/app"
+
+import streamlit as st
+import torch
+from torchvision import transforms
+from PIL import Image
+from TumorModel import TumorClassification, GliomaStageModel
+
+# 🎯 Load classification model
+tumor_model = TumorClassification()
+tumor_model.load_state_dict(torch.load("BTD_model.pth", map_location=torch.device("cpu")))
+tumor_model.eval()
+
+# 🎯 Load glioma stage model
+glioma_model = GliomaStageModel()
+glioma_model.load_state_dict(torch.load("glioma_stages.pth", map_location=torch.device("cpu")))
+glioma_model.eval()
+
+# 📌 Class labels
+tumor_labels = ['glioma', 'meningioma', 'notumor', 'pituitary']
+stage_labels = ['Stage 1', 'Stage 2', 'Stage 3', 'Stage 4']
+
+# 🔁 Image Transform
+transform = transforms.Compose([
+    transforms.Grayscale(),
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.5], std=[0.5])
+])
+
+# 🧠 Tumor prediction
+def predict_tumor(image: Image.Image) -> str:
+    image = transform(image).unsqueeze(0)
+    with torch.no_grad():
+        output = tumor_model(image)
+        pred = torch.argmax(output, dim=1).item()
+    return tumor_labels[pred]
+
+# 🧬 Glioma stage prediction
+def predict_stage(gender, age, idh1, tp53, atrx, pten, egfr, cic, pik3ca):
+    gender_val = 0 if gender == "Male" else 1
+    features = [gender_val, age, idh1, tp53, atrx, pten, egfr, cic, pik3ca]
+    x = torch.tensor(features).float().unsqueeze(0)
+    with torch.no_grad():
+        out = glioma_model(x)
+        pred = torch.argmax(out, dim=1).item()
+    return stage_labels[pred]
+
+# 🎨 Streamlit UI
+st.title("🧠 Brain Tumor Detection & Glioma Stage Predictor")
+
+tab1, tab2 = st.tabs(["Tumor Type Detector", "Glioma Stage Predictor"])
+
+with tab1:
+    st.header("🖼️ Upload an MRI Image")
+    uploaded_file = st.file_uploader("Choose an MRI image", type=["jpg", "jpeg", "png"])
+    if uploaded_file:
+        image = Image.open(uploaded_file).convert("RGB")
+        st.image(image, caption="Uploaded MRI", use_column_width=True)
+        prediction = predict_tumor(image)
+        st.success(f"🧠 Predicted Tumor Type: **{prediction.upper()}**")
+
+with tab2:
+    st.header("🧬 Patient Genetic and Demographic Information")
+    gender = st.radio("Gender", ["Male", "Female"])
+    age = st.slider("Age", 1, 100, 30)
+    idh1 = st.slider("IDH1 Mutation", 0, 1, 0)
+    tp53 = st.slider("TP53 Mutation", 0, 1, 0)
+    atrx = st.slider("ATRX Mutation", 0, 1, 0)
+    pten = st.slider("PTEN Mutation", 0, 1, 0)
+    egfr = st.slider("EGFR Mutation", 0, 1, 0)
+    cic = st.slider("CIC Mutation", 0, 1, 0)
+    pik3ca = st.slider("PIK3CA Mutation", 0, 1, 0)
+
+    if st.button("Predict Glioma Stage"):
+        stage = predict_stage(gender, age, idh1, tp53, atrx, pten, egfr, cic, pik3ca)
+        st.success(f"📊 Predicted Glioma Stage: **{stage}**")