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]}")