Create main.py

main.py

@@ -0,0 +1,43 @@
+import torch
+from transformers import AutoModel
+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_pred12')
+
+# 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)
+        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')