Spaces:
Sleeping
Sleeping
File size: 1,511 Bytes
e3e1e9b 3864676 e3e1e9b 862a52f e3e1e9b 3864676 e3e1e9b |
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 |
import torch
from transformers import AutoModel
import torch.nn as nn
from PIL import Image
import numpy as np
import streamlit as st
# Set the device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Load the trained model from the Hugging Face Hub
model = AutoModel.from_pretrained('dhhd255/parkinsons_pred0.1')
# Move the model to the device
model = model.to(device)
# Use Streamlit to upload an image
uploaded_file = st.file_uploader("Choose an image...", type=["png", "jpg", "jpeg"])
if uploaded_file is not None:
# Load and resize the image
image_size = (224, 224)
new_image = Image.open(uploaded_file).convert('RGB').resize(image_size)
new_image = np.array(new_image)
new_image = torch.from_numpy(new_image).transpose(0, 2).float().unsqueeze(0)
# Move the data to the device
new_image = new_image.to(device)
# Make predictions using the trained model
with torch.no_grad():
predictions = model(new_image)
logits = predictions.last_hidden_state
logits = logits.view(logits.shape[0], -1)
num_classes=2
feature_reducer = nn.Linear(logits.shape[1], num_classes)
logits = logits.to(device)
feature_reducer = feature_reducer.to(device)
logits = feature_reducer(logits)
predicted_class = torch.argmax(logits, dim=1).item()
if(predicted_class == 0):
st.write('Predicted class: Parkinson\'s')
else:
st.write('Predicted class: Healthy')
|