Update app.py

app.py

@@ -5,23 +5,30 @@ from PIL import Image
 
 # Authenticate and download the EfficientNet model from Hugging Face Spaces
 fs = HfFileSystem()
-model_path = 'dhhd255/efficientnet_b3/efficientnet_b3.pt'
-with fs.open(model_path, 'rb') as f:
-    model_content = f.read()
+efficientnet_model_path = 'dhhd255/efficientnet_b3/efficientnet_b3.pt'
+with fs.open(efficientnet_model_path, 'rb') as f:
+    efficientnet_model_content = f.read()
 
 # Save the EfficientNet model file to disk
 efficientnet_model_file = 'efficientnet_b3.pt'
 with open(efficientnet_model_file, 'wb') as f:
-    f.write(model_content)
+    f.write(efficientnet_model_content)
 
-# Load the EfficientNet model onto the CPU
-model = EfficientNet.from_pretrained('efficientnet-b3')
-model.load_state_dict(torch.load(efficientnet_model_file, map_location=torch.device('cpu')))
+# Authenticate and download your custom model from Hugging Face Spaces
+custom_model_path = 'dhhd255/efficient_net_parkinsons/best_model.pth'
+with fs.open(custom_model_path, 'rb') as f:
+    custom_model_content = f.read()
 
-# Load your custom model onto the CPU
+# Save your custom model file to disk
 custom_model_file = 'best_model.pth'
-model.load_state_dict(torch.load(custom_model_file, map_location=torch.device('cpu')))
+with open(custom_model_file, 'wb') as f:
+    f.write(custom_model_content)
 
+# Load the EfficientNet model onto the CPU
+model = torch.load(efficientnet_model_file)
+
+# Load your custom model onto the CPU
+model.load_state_dict(torch.load(custom_model_file))
 model.eval()
 
 # Define a function that takes an image as input and uses the model for inference
@@ -48,16 +55,16 @@ def image_classifier(image):
     # Return the result
     return outputs[0].numpy(), predicted_label
 
-# Create a Streamlit app with an image upload input
-uploaded_file = st.file_uploader('Upload an image')
-if uploaded_file is not None:
-    # Convert the UploadedFile object to a NumPy array
-    image = Image.open(uploaded_file)
-    image = np.array(image)
-    
-    # Use the image for inference
-    predictions, predicted_label = image_classifier(image)
-    
-    # Display the result
-    st.write(f'Predictions: {predictions}')
-    st.write(f'Predicted label: {predicted_label}')
+# Load and preprocess the image
+img_path = '/content/test_image_healthy.png'
+img = Image.open(img_path)
+img = data_transform(img)
+
+# Add a batch dimension and move the image to the device
+img = img.unsqueeze(0).to(device)
+
+# Perform inference
+with torch.no_grad():
+ outputs = model(img)
+ _, predicted = torch.max(outputs.data, 1)
+ print(f'Predicted class: {predicted.item()}')