Spaces:
Build error
Build error
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) | |
# Add custom CSS to use the Inter font, define custom classes for healthy and parkinsons results, increase the font size, make the text bold, and define the footer styles | |
st.markdown(""" | |
<style> | |
@import url('https://fonts.googleapis.com/css2?family=Inter&display=swap'); | |
body { | |
font-family: 'Inter', sans-serif; | |
} | |
.result { | |
font-size: 24px; | |
font-weight: bold; | |
} | |
.healthy { | |
color: #007E3F; | |
} | |
.parkinsons { | |
color: #C30000; | |
} | |
.caption_c{ | |
position: relative; | |
display: flex; | |
flex-directon: column; | |
align-items: center; | |
top: calc(99vh - 370px); | |
} | |
.caption { | |
text-align: center; | |
color: #646464; | |
font-size: 14px; | |
} | |
button:hover { | |
background-color: lightblue !important; | |
outline-color: lightblue !important; | |
} | |
button:focus { | |
background-color: lightblue !important; | |
outline-color: lightblue !important; | |
} | |
</style> | |
""", unsafe_allow_html=True) | |
st.title("Parkinson's Disease Prediction") | |
uploaded_file = st.file_uploader("Upload your :blue[Spiral] drawing here", type=["png", "jpg", "jpeg"]) | |
if uploaded_file is not None: | |
col1, col2 = st.columns(2) | |
# Load and resize the image | |
image_size = (224, 224) | |
new_image = Image.open(uploaded_file).convert('RGB').resize(image_size) | |
col1.image(new_image, use_column_width=True) | |
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() | |
confidence = torch.softmax(logits, dim=1)[0][predicted_class].item() | |
if(predicted_class == 0): | |
col2.markdown('<span class="result parkinsons">Predicted class: Parkinson\'s</span>', unsafe_allow_html=True) | |
col2.caption(f'{confidence*100:.0f}% sure') | |
else: | |
col2.markdown('<span class="result healthy">Predicted class: Healthy</span>', unsafe_allow_html=True) | |
col2.caption(f'{confidence*100:.0f}% sure') | |
# Add a caption at the bottom of the page | |
st.markdown('<div class="caption_c"><p class="caption">Made with love by Jayant</p></div>', unsafe_allow_html=True) |