Update app.py

app.py

@@ -1,70 +1,43 @@
 import torch
 from torchvision import transforms
-from huggingface_hub import HfFileSystem
+from huggingface_hub import HfApi, HfFolder
 from PIL import Image
 
-# Authenticate and download the EfficientNet model from Hugging Face Spaces
-fs = HfFileSystem()
-efficientnet_model_path = 'dhhd255/efficientnet_b3/efficientnet_b3.pt'
-with fs.open(efficientnet_model_path, 'rb') as f:
-    efficientnet_model_content = f.read()
+# Authenticate and download the EfficientNet model from Hugging Face
+api = HfApi()
+efficientnet_model_url = api.presigned_url('dhhd255/efficientnet_b3', filename='efficientnet_b3.pt').geturl()
+efficientnet_model_file = HfFolder.download_file(efficientnet_model_url, cache_dir=HfFolder.cache_dir())
 
-# Save the EfficientNet model file to disk
-efficientnet_model_file = 'efficientnet_b3.pt'
-with open(efficientnet_model_file, 'wb') as f:
-    f.write(efficientnet_model_content)
-
-# 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()
-
-# Save your custom model file to disk
-custom_model_file = 'best_model.pth'
-with open(custom_model_file, 'wb') as f:
-    f.write(custom_model_content)
+# Authenticate and download your custom model from Hugging Face
+custom_model_url = api.presigned_url('dhhd255/efficient_net_parkinsons', filename='best_model.pth').geturl()
+custom_model_file = HfFolder.download_file(custom_model_url, cache_dir=HfFolder.cache_dir())
 
 # Load the EfficientNet model onto the CPU
 model = torch.load(efficientnet_model_file, map_location=torch.device('cpu'))
 
-# Load your custom model onto the CPU
+# Load the saved weights onto the CPU
 model.load_state_dict(torch.load(custom_model_file, map_location=torch.device('cpu')))
+
+# Set the model to evaluation mode
 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 = data_transform(image)
-    image = image.unsqueeze(0)
-    
-    # Use your custom model for inference
-    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.item()]
-    
-    # Return the result
-    return outputs[0].numpy(), predicted_label
+# Define the image transform for inference
+data_transform = transforms.Compose([
+    transforms.Lambda(lambda x: x.convert('RGB')),
+    transforms.Resize((224, 224)),
+    transforms.ToTensor()
+])
 
 # 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)
+# Add a batch dimension
+img = img.unsqueeze(0)
 
 # Perform inference
 with torch.no_grad():
- outputs = model(img)
- _, predicted = torch.max(outputs.data, 1)
- print(f'Predicted class: {predicted.item()}')
+    outputs = model(img)
+    _, predicted = torch.max(outputs.data, 1)
+    print(f'Predicted class: {predicted.item()}')