Update app.py

app.py

@@ -9,13 +9,14 @@ app = Flask(__name__)
 # Load the trained model
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Define the model architecture
-model = models.resnet152()
-model.fc = torch.nn.Linear(model.fc.in_features, 26)  # Adjust for the number of classes
-model.load_state_dict(torch.load("trained_model.pth", map_location=device))
-model = model.to(device)
-model.eval()
-
+def load_model():
+    model = models.resnet152()
+    model.fc = torch.nn.Linear(model.fc.in_features, 26)  
+    model.load_state_dict(torch.load("trained_model.pth", map_location=device))
+    model = model.to(device)
+    model.eval()
+    return model
+    
 # Define preprocessing for the input image
 preprocess = transforms.Compose([
     transforms.Resize((224, 224)),
@@ -23,7 +24,7 @@ preprocess = transforms.Compose([
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 ])
 
-# Class labels (replace with your dataset's classes)
+# Class labels 
 CLASS_LABELS = [
 "bluebell", "buttercup", "colts_foot", "corn_poppy", "cowslip", 
 "crocus", "daffodil", "daisy", "dandelion", "foxglove", 
@@ -35,6 +36,7 @@ CLASS_LABELS = [
 
 @app.route("/predict", methods=["POST"])
 def predict():
+    model = load_model()
     if "file" not in request.files:
         return jsonify({"error": "No file uploaded"}), 400
 
@@ -56,6 +58,7 @@ def predict():
     except Exception as e:
         return jsonify({"error": f"Error during prediction: {str(e)}"}), 500
 
-# Run the app (Hugging Face Spaces requires `app.run()` here)
+# Run the app 
 if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=8080)
+    from waitress import serve  
+    serve(app, host="0.0.0.0", port=8080)