Update app.py

app.py

@@ -4,8 +4,9 @@ import torch.nn as nn
 from torchvision import transforms
 from PIL import Image
 import io
+from torch.hub import load_state_dict_from_url
 
-# Define the model architecture that matches the saved .pth file
+# Define the model architecture
 class BacterialMorphologyClassifier(nn.Module):
     def __init__(self):
         super(BacterialMorphologyClassifier, self).__init__()
@@ -31,16 +32,13 @@ class BacterialMorphologyClassifier(nn.Module):
         x = self.fc(x)
         return x
 
-# Load the model and weights at the start of the app
+# Load the model and weights
 model = BacterialMorphologyClassifier()
-MODEL_PATH = 'https://huggingface.co/yolac/BacterialMorphologyClassification/resolve/main/model.pth'  # Replace this with the local path if needed
-model.load_state_dict(torch.load(MODEL_PATH, map_location=torch.device('cpu')), strict=False)
+MODEL_URL = "https://huggingface.co/yolac/BacterialMorphologyClassification/resolve/main/model.pth"
+state_dict = load_state_dict_from_url(MODEL_URL, map_location=torch.device('cpu'))
+model.load_state_dict(state_dict, strict=False)
 model.eval()
 
-# Move model to GPU if available
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model.to(device)
-
 # Set up Flask app
 app = Flask(__name__)
 
@@ -59,7 +57,7 @@ def predict():
         image = Image.open(io.BytesIO(image_file.read())).convert('RGB')
         
         # Preprocess the image
-        image_tensor = transform(image).unsqueeze(0).to(device)
+        image_tensor = transform(image).unsqueeze(0)
         
         # Make prediction
         output = model(image_tensor)
@@ -78,4 +76,4 @@ def predict():
         return jsonify({'error': str(e)})
 
 if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=5000, debug=False)  # Set debug=False for production
+    app.run(host='0.0.0.0', port=5000, debug=False)