Update app.py

app.py

@@ -38,11 +38,14 @@ class BacterialMorphologyClassifier(nn.Module):
 # Load the model and weights
 MODEL_PATH = "https://huggingface.co/yolac/BacterialMorphologyClassification/resolve/main/model.pth"
 logging.debug("Starting model loading...")
-model = BacterialMorphologyClassifier()
-state_dict = torch.hub.load_state_dict_from_url(MODEL_PATH, map_location=torch.device('cpu'))
-model.load_state_dict(state_dict, strict=False)
-model.eval()
-logging.debug("Model loaded successfully.")
+try:
+    model = BacterialMorphologyClassifier()
+    state_dict = torch.hub.load_state_dict_from_url(MODEL_PATH, map_location=torch.device('cpu'))
+    model.load_state_dict(state_dict, strict=False)
+    model.eval()
+    logging.debug("Model loaded successfully.")
+except Exception as e:
+    logging.error(f"Error loading the model: {str(e)}")
 
 # Define image preprocessing transformations
 transform = transforms.Compose([
@@ -54,8 +57,10 @@ transform = transforms.Compose([
 # Define the prediction function
 def predict(image):
     try:
+        logging.debug("Starting prediction...")
         # Preprocess the image
         image_tensor = transform(image).unsqueeze(0)
+        logging.debug("Image preprocessing completed.")
         
         # Make prediction
         output = model(image_tensor)