app.py

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}**")