import streamlit as st
import torch
import torch.nn as nn
from torchvision import transforms, datasets, models
from PIL import Image
st.markdown(
"""
""",
unsafe_allow_html=True
)
# Title
st.title("Brain Tumor Classification")
st.write("")
# Class names
class_names = ['glioma_tumor', 'meningioma_tumor', 'no_tumor', 'pituitary_tumor']
# Load pre-trained ResNet18 model
model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
num_of_classes = len(class_names)
num_of_features = model.fc.in_features
model.fc = nn.Linear(num_of_features, num_of_classes)
# Load trained model weights
model.load_state_dict(torch.load('resnet18_model (1).pth', map_location=torch.device('cpu')))
model.eval()
st.markdown(
"""
📤 Upload a Scan Image
""",
unsafe_allow_html=True
)
# Image upload
uploaded_img = st.file_uploader("", type=["jpg", "jpeg", "png"])
if st.button("Submit"):
if uploaded_img is not None:
# Display uploaded image in a smaller size
image = Image.open(uploaded_img)
st.image(image, caption="**Uploaded Image**", width=200)
# Image transformations
sample_transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.1776, 0.1776, 0.1776], std=[0.1735, 0.1735, 0.1735])
])
# Apply transformations
transformed_img = sample_transform(image).unsqueeze(0)
# Model inference
with torch.no_grad():
pred = model(transformed_img).argmax(dim=1).item()
# Stylish output box
st.markdown(
f"""
🧠 Predicted Class: {class_names[pred]}
""",
unsafe_allow_html=True
)
else:
st.markdown(
"""
⚠️ Plese upload image
""",
unsafe_allow_html=True
)