File size: 2,254 Bytes
9c1eddf |
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 |
import torch
from transformers import SegformerForImageClassification
from torchvision import transforms
from PIL import Image
import gradio as gr
# Load Alzheimer's model
alzheimers_model = SegformerForImageClassification.from_pretrained('nvidia/mit-b1')
alzheimers_model.classifier = torch.nn.Linear(alzheimers_model.classifier.in_features, 4) # 4 classes
alzheimers_model.load_state_dict(torch.load('alzheimers_model.pth', map_location=torch.device('cpu')))
alzheimers_model.eval()
# Load Brain Tumor model
brain_tumor_model = SegformerForImageClassification.from_pretrained('nvidia/mit-b1')
brain_tumor_model.classifier = torch.nn.Linear(brain_tumor_model.classifier.in_features, 4) # 4 classes
brain_tumor_model.load_state_dict(torch.load('brain_tumor_model.pth', map_location=torch.device('cpu')))
brain_tumor_model.eval()
# Define transformations
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
# Prediction function for Alzheimer's
def predict_alzheimers(image):
image = transform(image).unsqueeze(0)
with torch.no_grad():
outputs = alzheimers_model(image).logits
_, predicted = torch.max(outputs, 1)
classes = ['Mild Dementia', 'Moderate Dementia', 'Non Demented', 'Very mild Dementia']
return classes[predicted.item()]
# Prediction function for Brain Tumor
def predict_brain_tumor(image):
image = transform(image).unsqueeze(0)
with torch.no_grad():
outputs = brain_tumor_model(image).logits
_, predicted = torch.max(outputs, 1)
classes = ['glioma', 'meningioma', 'notumor', 'pituitary']
return classes[predicted.item()]
def predict(image, model_type):
if model_type == "Alzheimer's":
return predict_alzheimers(image)
elif model_type == "Brain Tumor":
return predict_brain_tumor(image)
interface = gr.Interface(
fn=predict,
inputs=[gr.Image(type="pil"), gr.Dropdown(["Alzheimer's", "Brain Tumor"])],
outputs=gr.Textbox(),
title="MRI Scan Classification",
description="Upload an MRI scan and select the type of classification."
)
interface.launch() |