Update app.py

app.py

@@ -1,43 +1,63 @@
+import io
 import torch
 from torchvision import transforms
-from huggingface_hub import HfApi, HfFolder
+from huggingface_hub import HfFileSystem
 from PIL import Image
 
 # 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())
-
-# 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())
+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()
 
 # Load the EfficientNet model onto the CPU
+efficientnet_model_file = io.BytesIO(efficientnet_model_content)
 model = torch.load(efficientnet_model_file, map_location=torch.device('cpu'))
 
-# Load the saved weights onto the CPU
-model.load_state_dict(torch.load(custom_model_file, map_location=torch.device('cpu')))
+# Authenticate and download your custom model from Hugging Face
+custom_model_path = 'dhhd255/efficient_net_parkinsons/best_model.pth'
+with fs.open(custom_model_path, 'rb') as f:
+    custom_model_content = f.read()
 
-# Set the model to evaluation mode
+# Load your custom model onto the CPU
+custom_model_file = io.BytesIO(custom_model_content)
+model.load_state_dict(torch.load(custom_model_file, map_location=torch.device('cpu')))
 model.eval()
 
-# Define the image transform for inference
-data_transform = transforms.Compose([
-    transforms.Lambda(lambda x: x.convert('RGB')),
-    transforms.Resize((224, 224)),
-    transforms.ToTensor()
-])
+# 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
 
 # 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
+# Add a batch dimension and move the image to the device
 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()}')