File size: 3,038 Bytes
b89d505 |
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 83 84 85 |
import streamlit as st
from ultralytics import YOLO
from PIL import Image
import torchvision.transforms as transforms
import base64
# Set Streamlit Page Configuration
st.set_page_config(
page_title="Brain Tumor Segmentation",
page_icon="logo/logo.png",
layout="centered"
)
# Cache the YOLO model to avoid reloading on every interaction
@st.cache_resource()
def load_model():
return YOLO("model/best.pt") # Update path if needed
model = load_model()
# Define image transformation pipeline
transform = transforms.Compose([
transforms.Resize((640, 640)),
transforms.ToTensor()
])
# Function to predict and overlay tumor segmentation mask
def predict_tumor(image: Image.Image):
try:
image_tensor = transform(image).unsqueeze(0) # Add batch dimension
results = model.predict(image_tensor)
output_image = results[0].plot() # Overlay segmentation mask
return Image.fromarray(output_image)
except Exception as e:
st.error(f"Prediction Error: {e}")
return None
# Function to encode image to base64 for embedding
def get_base64_image(image_path):
with open(image_path, "rb") as img_file:
return base64.b64encode(img_file.read()).decode()
# Display logo
image_base64 = get_base64_image("logo/logo.png")
st.markdown(
f'<div style="text-align: center;"><img src="data:image/png;base64,{image_base64}" width="100"></div>',
unsafe_allow_html=True
)
# --- UI Customization ---
st.markdown("""
<style>
[data-testid="stSidebar"] { background-color: #1E1E2F; }
[data-testid="stSidebar"] h1, [data-testid="stSidebar"] h2 { color: white; }
h1 { text-align: center; font-size: 36px; font-weight: bold; color: #2C3E50; }
div.stButton > button { background-color: #3498DB; color: white; font-weight: bold; }
div.stButton > button:hover { background-color: #2980B9; }
</style>
""", unsafe_allow_html=True)
# --- Sidebar ---
st.sidebar.header("📤 Upload an MRI Image")
uploaded_file = st.sidebar.file_uploader("Drag and drop or browse", type=['jpg', 'png', 'jpeg'])
# --- Main Page ---
st.title("Brain Tumor Segmentation")
st.markdown("<p style='text-align: center;'>Detect and segment brain tumors from MRI scans.</p>", unsafe_allow_html=True)
if uploaded_file:
image = Image.open(uploaded_file).convert("RGB")
col1, col2 = st.columns(2)
with col1:
st.image(image, caption="📷 Uploaded Image", use_container_width=True)
if st.sidebar.button("🔍 Predict Tumor Segmentation"):
segmented_image = predict_tumor(image)
if segmented_image:
with col2:
st.image(segmented_image, caption="🎯 Segmented Tumor", use_container_width=True)
else:
st.error("Segmentation failed. Please try again.")
st.markdown("---")
st.info("This app uses **YOLO-Seg** for real-time tumor segmentation. Upload an MRI image to get started.") |