Update app.py

app.py

@@ -31,21 +31,20 @@ class BacterialMorphologyClassifier(nn.Module):
         x = self.fc(x)
         return x
 
-# Load the model and weights
+# Load the model and weights at app startup
 model = BacterialMorphologyClassifier()
 MODEL_PATH = "https://huggingface.co/yolac/BacterialMorphologyClassification/resolve/main/model.pth"
 state_dict = torch.hub.load_state_dict_from_url(MODEL_PATH, map_location=torch.device('cpu'))
-model.load_state_dict(state_dict)
+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__)
 
-# Add a basic route
-@app.route('/')
-def home():
-    return "Flask app is running!"
-
 # Define image preprocessing transformations
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
@@ -61,7 +60,7 @@ def predict():
         image = Image.open(io.BytesIO(image_file.read())).convert('RGB')
         
         # Preprocess the image
-        image_tensor = transform(image).unsqueeze(0)
+        image_tensor = transform(image).unsqueeze(0).to(device)
         
         # Make prediction
         output = model(image_tensor)