File size: 1,425 Bytes
b76c318 |
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 |
import streamlit as st
import torch
import torch.nn as nn
from torchvision import transforms, datasets, models
from PIL import Image
# Title
st.title("Brain Tumor Classification")
# 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()
# Image upload
uploaded_img = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
if uploaded_img is not None:
# Display uploaded image
image = Image.open(uploaded_img)
st.image(image, caption="Uploaded Image", use_container_width =True)
# 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()
# Display prediction
st.success(f"Predicted Class: {class_names[pred]}")
|