Spaces:
Sleeping
Sleeping
File size: 2,952 Bytes
5760faf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
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}**")
|