Spaces:
Sleeping
Sleeping
import gradio as gr | |
import torch | |
from torchvision import transforms | |
from PIL import Image | |
from model import load_model | |
import numpy as np | |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
model = load_model(device) | |
class_names = [ | |
'Alzheimer Disease', | |
'Mild Alzheimer Risk', | |
'Moderate Alzheimer Risk', | |
'Very Mild Alzheimer Risk', | |
'No Risk', | |
'Parkinson Disease' | |
] | |
transform = transforms.Compose([ | |
transforms.Resize((448, 448)), | |
transforms.ToTensor(), | |
transforms.Normalize(mean=[0.485, 0.456, 0.406], | |
std=[0.229, 0.224, 0.225]) | |
]) | |
def predict(image): | |
image = image.convert("RGB") | |
tensor = transform(image).unsqueeze(0).to(device) | |
with torch.no_grad(): | |
outputs = model(tensor) | |
probs = torch.nn.functional.softmax(outputs, dim=1)[0] | |
predicted = torch.argmax(probs).item() | |
confidence = probs[predicted].item() * 100 | |
return {class_names[i]: float(probs[i]) for i in range(len(class_names))} | |
gr.Interface( | |
fn=predict, | |
inputs=gr.Image(type="pil"), | |
outputs=gr.Label(num_top_classes=3), | |
title="Vbai-DPA 2.2 (C Version)", | |
description="Upload an MRI and fMRI image to classify the risk level using the 'C' version of the Vbai-DPA 2.2 model." | |
).launch() | |