Update app.py

app.py

@@ -1,45 +1,46 @@
 import streamlit as st
-from tensorflow import keras
-import numpy as np
+import torch
+from torchvision import transforms
 from huggingface_hub import HfFileSystem
 from PIL import Image
 
 # Authenticate and download the custom model from Hugging Face Spaces
 fs = HfFileSystem()
-model_path = 'dhhd255/main_model/best_model.h5'
+model_path = 'dhhd255/main_model/best_model.pth'
 with fs.open(model_path, 'rb') as f:
     model_content = f.read()
 
 # Save the model file to disk
-with open('best_model.h5', 'wb') as f:
+with open('best_model.pth', 'wb') as f:
     f.write(model_content)
 
 # Load your custom model
-model = keras.models.load_model('best_model.h5')
+model = torch.load('best_model.pth')
+model.eval()
 
 # Define a function that takes an image as input and uses the model for inference
 def image_classifier(image):
     # Preprocess the input image
+    data_transform = transforms.Compose([
+        transforms.Lambda(lambda x: x.convert('RGB')),
+        transforms.Resize((224, 224)),
+        transforms.ToTensor()
+    ])
     image = Image.fromarray(image)
-    image = image.convert('L')
-    image = image.resize((128, 128))
-    image = np.array(image)
-    image = image / 255.0
-    image = np.expand_dims(image, axis=0)
-    image = np.expand_dims(image, axis=-1)
+    image = data_transform(image)
+    image = image.unsqueeze(0)
     
     # Use your custom model for inference
-    predictions = model.predict(image)
-    
-    # Get the index of the highest predicted probability
-    predicted_index = np.argmax(predictions[0])
+    with torch.no_grad():
+        outputs = model(image)
+        _, predicted = torch.max(outputs.data, 1)
     
     # Map the index to a class label
     labels = ['Healthy', 'Parkinson']
-    predicted_label = labels[predicted_index]
+    predicted_label = labels[predicted.item()]
     
     # Return the result
-    return predictions[0], predicted_label
+    return outputs[0].numpy(), predicted_label
 
 # Create a Streamlit app with an image upload input
 uploaded_file = st.file_uploader('Upload an image')