Update app.py

app.py

@@ -4,19 +4,25 @@ from efficientnet_pytorch import EfficientNet
 from huggingface_hub import HfFileSystem
 from PIL import Image
 
-# Authenticate and download the custom model from Hugging Face Spaces
+# 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()
 
-# Save the model file to disk
-with open('efficientnet_b3.pt', 'wb') as f:
+# 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)
 
-# Load your custom model onto the CPU
+# Load the EfficientNet model onto the CPU
 model = EfficientNet.from_pretrained('efficientnet-b3')
-model.load_state_dict(torch.load('efficientnet_b3.pt', map_location=torch.device('cpu')))
+model.load_state_dict(torch.load(efficientnet_model_file, map_location=torch.device('cpu')))
+
+# Load your custom model onto the CPU
+custom_model_file = 'best_model.pth'
+model.load_state_dict(torch.load(custom_model_file, map_location=torch.device('cpu')))
+
 model.eval()
 
 # Define a function that takes an image as input and uses the model for inference