Bug : model.py & utils.py – Unhandled Exceptions on Model Load

Disease prediction/model.py

@@ -1,4 +1,9 @@
+import torch
 import torch.nn as nn
+import os
+import logging
+
+logger = logging.getLogger(__name__)
 
 class PlantDiseaseNet(nn.Module):
     def __init__(self, num_classes=38):  # original model had 38 classes
@@ -77,3 +82,34 @@ def forward(self, x):
         x = self.pool(x)
         x = self.classifier(x)
         return x
+
+def load_pytorch_model(model_path, device='cpu'):
+    try:
+        if not os.path.exists(model_path):
+            logger.error(f"PyTorch model file missing: {model_path}")
+            return None
+        # Safe loading with torch.load
+        model = torch.load(model_path, map_location=device)
+        # If it was saved as state dict rather than full model
+        if isinstance(model, dict):
+            net = PlantDiseaseNet()
+            net.load_state_dict(model)
+            model = net
+        model.eval()
+        return model
+    except Exception as e:
+        logger.error(f"Error loading PyTorch model from {model_path}: {str(e)}")
+        return None
+
+def predict_pytorch(model, input_tensor):
+    if model is None:
+        logger.error("Prediction failed: PyTorch model is not loaded (None)")
+        return None
+    try:
+        with torch.no_grad():
+            output = model(input_tensor)
+            return output
+    except Exception as e:
+        logger.error(f"Error predicting with PyTorch model: {str(e)}")
+        return None
+

Disease prediction/utils.py

@@ -3,6 +3,9 @@
 from tensorflow.keras.models import load_model
 from tensorflow.keras.preprocessing import image
 import os
+import logging
+
+logger = logging.getLogger(__name__)
 
 # Class labels (update if needed)
 class_names = [
@@ -42,19 +45,34 @@
 }
 
 def load_keras_model(model_path):
-    model = load_model(model_path)
-    return model
-
-
+    try:
+        if not os.path.exists(model_path):
+            logger.error(f"Model file does not exist at path: {model_path}")
+            return None
+        model = load_model(model_path)
+        return model
+    except Exception as e:
+        logger.error(f"Failed to load Keras model from {model_path}: {str(e)}")
+        return None
 
 def predict_image_keras(model, img_path):
-    img = image.load_img(img_path, target_size=(160, 160))  # Match your training size
-    img_array = image.img_to_array(img) / 255.0
-    img_array = np.expand_dims(img_array, axis=0)
+    if model is None:
+        logger.error("Prediction failed: model is None")
+        return "Model unavailable", "The classification model is currently unavailable."
+    try:
+        if not os.path.exists(img_path):
+            logger.error(f"Image file does not exist at path: {img_path}")
+            return "Image not found", "The uploaded image file could not be found."
+        img = image.load_img(img_path, target_size=(160, 160))  # Match your training size
+        img_array = image.img_to_array(img) / 255.0
+        img_array = np.expand_dims(img_array, axis=0)
 
-    predictions = model.predict(img_array)
-    predicted_index = np.argmax(predictions)
-    predicted_class = class_names[predicted_index]
-    description = class_descriptions.get(predicted_class, "No description available.")
+        predictions = model.predict(img_array)
+        predicted_index = np.argmax(predictions)
+        predicted_class = class_names[predicted_index]
+        description = class_descriptions.get(predicted_class, "No description available.")
 
-    return predicted_class, description
+        return predicted_class, description
+    except Exception as e:
+        logger.error(f"Error predicting image {img_path}: {str(e)}")
+        return "Model unavailable", f"An error occurred during prediction: {str(e)}"

model.py

@@ -1,4 +1,9 @@
+import torch
 import torch.nn as nn
+import os
+import logging
+
+logger = logging.getLogger(__name__)
 
 class PlantDiseaseNet(nn.Module):
     def __init__(self, num_classes=38):  # original model had 38 classes
@@ -77,3 +82,34 @@ def forward(self, x):
         x = self.pool(x)
         x = self.classifier(x)
         return x
+
+def load_pytorch_model(model_path, device='cpu'):
+    try:
+        if not os.path.exists(model_path):
+            logger.error(f"PyTorch model file missing: {model_path}")
+            return None
+        # Safe loading with torch.load
+        model = torch.load(model_path, map_location=device)
+        # If it was saved as state dict rather than full model
+        if isinstance(model, dict):
+            net = PlantDiseaseNet()
+            net.load_state_dict(model)
+            model = net
+        model.eval()
+        return model
+    except Exception as e:
+        logger.error(f"Error loading PyTorch model from {model_path}: {str(e)}")
+        return None
+
+def predict_pytorch(model, input_tensor):
+    if model is None:
+        logger.error("Prediction failed: PyTorch model is not loaded (None)")
+        return None
+    try:
+        with torch.no_grad():
+            output = model(input_tensor)
+            return output
+    except Exception as e:
+        logger.error(f"Error predicting with PyTorch model: {str(e)}")
+        return None
+

tests/conftest.py

@@ -8,11 +8,15 @@
 # Provide lightweight stubs for optional heavy external modules so tests can import app
 import types
 if 'google.generativeai' not in sys.modules:
-    google = types.ModuleType('google')
+    if 'google' in sys.modules:
+        google = sys.modules['google']
+    else:
+        google = types.ModuleType('google')
+        google.__path__ = []
+        sys.modules['google'] = google
     genai = types.ModuleType('google.generativeai')
-    google.generativeai = genai
-    sys.modules['google'] = google
     sys.modules['google.generativeai'] = genai
+    setattr(google, 'generativeai', genai)
 
 # Import the Flask app from the project
 try:

tests/test_predictions.py

@@ -48,3 +48,41 @@ def test_disease_predict_without_file(self, client):
         response = client.post('/disease/predict')
         # Should redirect to home when no file
         assert response.status_code in [200, 302, 400]
+
+
+class TestMLModelErrorHandling:
+    """Test suite for ML model load and prediction error handling."""
+
+    def test_load_keras_model_missing_file(self):
+        """Test that load_keras_model returns None when file is missing."""
+        from utils import load_keras_model
+        assert load_keras_model("nonexistent_model.h5") is None
+
+    def test_predict_image_keras_none_model(self):
+        """Test that predict_image_keras returns controlled error response when model is None."""
+        from utils import predict_image_keras
+        pred, desc = predict_image_keras(None, "nonexistent_image.jpg")
+        assert pred == "Model unavailable"
+        assert "unavailable" in desc
+
+    def test_predict_image_keras_missing_image(self):
+        """Test that predict_image_keras handles missing image files gracefully when model is mocked."""
+        from utils import predict_image_keras
+        class MockModel:
+            def predict(self, x):
+                return [[0.9, 0.1]]
+
+        pred, desc = predict_image_keras(MockModel(), "nonexistent_image.jpg")
+        assert pred == "Image not found"
+        assert "could not be found" in desc
+
+    def test_load_pytorch_model_missing_file(self):
+        """Test that load_pytorch_model returns None when file is missing."""
+        from model import load_pytorch_model
+        assert load_pytorch_model("nonexistent_model.pth") is None
+
+    def test_predict_pytorch_none_model(self):
+        """Test that predict_pytorch returns None when model is None."""
+        from model import predict_pytorch
+        assert predict_pytorch(None, None) is None
+

utils.py

@@ -3,6 +3,9 @@
 from tensorflow.keras.models import load_model
 from tensorflow.keras.preprocessing import image
 import os
+import logging
+
+logger = logging.getLogger(__name__)
 
 # Class labels (update if needed)
 class_names = [
@@ -42,19 +45,34 @@
 }
 
 def load_keras_model(model_path):
-    model = load_model(model_path)
-    return model
-
-
+    try:
+        if not os.path.exists(model_path):
+            logger.error(f"Model file does not exist at path: {model_path}")
+            return None
+        model = load_model(model_path)
+        return model
+    except Exception as e:
+        logger.error(f"Failed to load Keras model from {model_path}: {str(e)}")
+        return None
 
 def predict_image_keras(model, img_path):
-    img = image.load_img(img_path, target_size=(160, 160))  # Match your training size
-    img_array = image.img_to_array(img) / 255.0
-    img_array = np.expand_dims(img_array, axis=0)
+    if model is None:
+        logger.error("Prediction failed: model is None")
+        return "Model unavailable", "The classification model is currently unavailable."
+    try:
+        if not os.path.exists(img_path):
+            logger.error(f"Image file does not exist at path: {img_path}")
+            return "Image not found", "The uploaded image file could not be found."
+        img = image.load_img(img_path, target_size=(160, 160))  # Match your training size
+        img_array = image.img_to_array(img) / 255.0
+        img_array = np.expand_dims(img_array, axis=0)
 
-    predictions = model.predict(img_array)
-    predicted_index = np.argmax(predictions)
-    predicted_class = class_names[predicted_index]
-    description = class_descriptions.get(predicted_class, "No description available.")
+        predictions = model.predict(img_array)
+        predicted_index = np.argmax(predictions)
+        predicted_class = class_names[predicted_index]
+        description = class_descriptions.get(predicted_class, "No description available.")
 
-    return predicted_class, description
+        return predicted_class, description
+    except Exception as e:
+        logger.error(f"Error predicting image {img_path}: {str(e)}")
+        return "Model unavailable", f"An error occurred during prediction: {str(e)}"