SURESHBEEKHANI/Brain_Tumor_Segmentation

1import streamlit as st
2from ultralytics import YOLO
3from PIL import Image
4import torchvision.transforms as transforms
5import base64
6
7# Set Streamlit Page Configuration
8st.set_page_config(
9    page_title="Brain Tumor Segmentation",
10    page_icon="logo/logo.png",
11    layout="centered"
12)
13
14# Cache the YOLO model to avoid reloading on every interaction
15@st.cache_resource()
16def load_model():
17    return YOLO("model/best.pt")  # Update path if needed
18
19model = load_model()
20
21# Define image transformation pipeline
22transform = transforms.Compose([
23    transforms.Resize((640, 640)),
24    transforms.ToTensor()
25])
26
27# Function to predict and overlay tumor segmentation mask
28def predict_tumor(image: Image.Image):
29    try:
30        image_tensor = transform(image).unsqueeze(0)  # Add batch dimension
31        results = model.predict(image_tensor)
32        output_image = results[0].plot()  # Overlay segmentation mask
33        return Image.fromarray(output_image)
34    except Exception as e:
35        st.error(f"Prediction Error: {e}")
36        return None
37
38# Function to encode image to base64 for embedding
39def get_base64_image(image_path):
40    with open(image_path, "rb") as img_file:
41        return base64.b64encode(img_file.read()).decode()
42
43# Display logo
44image_base64 = get_base64_image("logo/logo.png")
45st.markdown(
46    f'<div style="text-align: center;"><img src="data:image/png;base64,{image_base64}" width="100"></div>',
47    unsafe_allow_html=True
48)
49
50# --- UI Customization ---
51st.markdown("""
52    <style>
53        [data-testid="stSidebar"] { background-color: #1E1E2F; }
54        [data-testid="stSidebar"] h1, [data-testid="stSidebar"] h2 { color: white; }
55        h1 { text-align: center; font-size: 36px; font-weight: bold; color: #2C3E50; }
56        div.stButton > button { background-color: #3498DB; color: white; font-weight: bold; }
57        div.stButton > button:hover { background-color: #2980B9; }
58    </style>
59""", unsafe_allow_html=True)
60
61# --- Sidebar ---
62st.sidebar.header("📤 Upload an MRI Image")
63uploaded_file = st.sidebar.file_uploader("Drag and drop or browse", type=['jpg', 'png', 'jpeg'])
64
65# --- Main Page ---
66st.title("Brain Tumor Segmentation")
67st.markdown("<p style='text-align: center;'>Detect and segment brain tumors from MRI scans.</p>", unsafe_allow_html=True)
68
69if uploaded_file:
70    image = Image.open(uploaded_file).convert("RGB")
71    col1, col2 = st.columns(2)
72    
73    with col1:
74        st.image(image, caption="📷 Uploaded Image", use_container_width=True)
75    
76    if st.sidebar.button("🔍 Predict Tumor Segmentation"):
77        segmented_image = predict_tumor(image)
78        if segmented_image:
79            with col2:
80                st.image(segmented_image, caption="🎯 Segmented Tumor", use_container_width=True)
81        else:
82            st.error("Segmentation failed. Please try again.")
83
84st.markdown("---")
85st.info("This app uses **YOLO-Seg** for real-time tumor segmentation. Upload an MRI image to get started.")