Update app.py

app.py

@@ -33,7 +33,7 @@ class BacterialMorphologyClassifier(nn.Module):
         return x
 
 # Load the model
-MODEL_PATH = "model.pth"
+MODEL_PATH = "https://huggingface.co/yolac/BacterialMorphologyClassification/resolve/main/model.pth"
 model = BacterialMorphologyClassifier()
 
 try:
@@ -45,7 +45,7 @@ try:
         with open(MODEL_PATH, "wb") as f:
             f.write(response.content)
     # Load the model weights
-    model.load_state_dict(torch.load(MODEL_PATH, map_location=torch.device('cpu')), strict=False)
+    model.load_state_dict(torch.load(MODEL_PATH, map_location=torch.device('cpu')))
     model.eval()
     print("Model loaded successfully.")
 except Exception as e:
@@ -65,15 +65,15 @@ def predict(image):
         image_tensor = transform(image).unsqueeze(0)
         
         # Perform prediction
-        output = model(image_tensor)
-        prediction = output.argmax().item()
+        with torch.no_grad():  # Ensure no gradients are calculated
+            output = model(image_tensor)
         
         # Class mapping
         class_labels = {0: 'cocci', 1: 'bacilli', 2: 'spirilla'}
         
         # Return the predicted class and confidence
-        predicted_class = class_labels[prediction]
-        confidence = output.max().item()
+        predicted_class = class_labels[output.argmax().item()]
+        confidence = output.max().item()  # Softmax value as confidence
         return f"Predicted Class: {predicted_class}\nConfidence: {confidence:.2f}"
     except Exception as e:
         return f"Error: {str(e)}"