Update main.py

main.py

@@ -8,13 +8,16 @@ from tensorflow.keras.preprocessing import image
 from tensorflow.keras.layers import Layer, Conv2D, Softmax, Concatenate
 import shutil
 import uvicorn
-import requests
 
 app = FastAPI()
 
 # Directory where models are stored
 MODEL_DIRECTORY = "dsanet_models"
 
+# Temporary directory for uploaded files
+TMP_DIR = os.getenv("TMP_DIR", "/app/temp")
+os.makedirs(TMP_DIR, exist_ok=True)  # Ensure the temp directory exists
+
 # Plant disease class names
 plant_disease_dict = {
     "Rice": ['Blight', 'Brown_Spots'],
@@ -34,8 +37,7 @@ plant_disease_dict = {
     "Corn": ['Corn___Cercospora_leaf_spot Gray_leaf_spot', 'Corn___Common_rust',
              'Corn___Northern_Leaf_Blight', 'Corn___healthy']
 }
-TMP_DIR = os.getenv("TMP_DIR", "/app/temp")
-os.makedirs(TMP_DIR, exist_ok=True)
+
 # Custom Self-Attention Layer
 @tf.keras.utils.register_keras_serializable()
 class SelfAttention(Layer):
@@ -69,9 +71,39 @@ class SelfAttention(Layer):
         config.update({"reduction_ratio": self.reduction_ratio})
         return config
 
+
+# **Load all models into memory at startup**
+loaded_models = {}
+
+def load_all_models():
+    """
+    Load all models from the `dsanet_models` directory at startup.
+    """
+    global loaded_models
+    for plant_name in plant_disease_dict.keys():
+        model_path = os.path.join(MODEL_DIRECTORY, f"model_{plant_name}.keras")
+
+        if os.path.isfile(model_path):
+            try:
+                if plant_name == "Rice":
+                    loaded_models[plant_name] = load_model(model_path)  # Load normally
+                else:
+                    loaded_models[plant_name] = load_model(model_path, custom_objects={"SelfAttention": SelfAttention})
+                print(f"✅ Model for {plant_name} loaded successfully!")
+            except Exception as e:
+                print(f"❌ Error loading model '{plant_name}': {e}")
+        else:
+            print(f"⚠ Warning: Model file '{model_path}' not found!")
+
+# Load models at startup
+load_all_models()
+
+
 @app.get("/health")
 async def api_health_check():
     return JSONResponse(content={"status": "Service is running"})
+
+
 @app.post("/predict/{plant_name}")
 async def predict_plant_disease(plant_name: str, file: UploadFile = File(...)):
     """
@@ -85,32 +117,17 @@ async def predict_plant_disease(plant_name: str, file: UploadFile = File(...)):
         JSON response with the predicted class.
     """
     # Ensure the plant name is valid
-    if plant_name not in plant_disease_dict:
-        raise HTTPException(status_code=400, detail="Invalid plant name")
-
-    # Construct the model path
-    model_path = os.path.join(MODEL_DIRECTORY, f"model_{plant_name}.keras")
-    if plant_name == "Rice":
-        model = load_model(model_path)
-    else:
-        model = load_model(model_path, custom_objects={"SelfAttention": SelfAttention})
-
-
-    # Check if the model exists
-    if not os.path.isfile(model_path):
-        raise HTTPException(status_code=404, detail=f"Model file '{plant_name}_model.keras' not found")
+    if plant_name not in loaded_models:
+        raise HTTPException(status_code=400, detail=f"Invalid plant name or model not loaded: {plant_name}")
 
     # Save uploaded file temporarily
-    
-
-    # Define the temp file path
     temp_path = os.path.join(TMP_DIR, file.filename)
     with open(temp_path, "wb") as buffer:
         shutil.copyfileobj(file.file, buffer)
 
     try:
-        # Load model
-        model = load_model(model_path, custom_objects={"SelfAttention": SelfAttention})
+        # Retrieve the preloaded model
+        model = loaded_models[plant_name]
 
         # Load and preprocess the image
         img = image.load_img(temp_path, target_size=(224, 224))
@@ -126,5 +143,7 @@ async def predict_plant_disease(plant_name: str, file: UploadFile = File(...)):
     finally:
         # Clean up temporary file
         os.remove(temp_path)
+
+
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)