Fixed Model Loading Issue for CPU-Only Environments

app.py

@@ -64,9 +64,9 @@ def visualize(image, mask, box):
 def process_images(file, x_min, y_min, x_max, y_max):
     image, H, W = load_image(file)
 
-    # Initialize MedSAM model
+    # Check if CUDA is available, and set the device accordingly
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
-    
+
     # Load the model with appropriate device mapping
     model_checkpoint_path = "medsam_vit_b.pth"  
     if device == 'cpu':
@@ -74,7 +74,12 @@ def process_images(file, x_min, y_min, x_max, y_max):
     else:
         checkpoint = torch.load(model_checkpoint_path)
 
-    medsam_model = sam_model_registry['vit_b'](checkpoint=checkpoint)
+    # Create model instance
+    medsam_model = sam_model_registry['vit_b']()  # Create model instance without checkpoint
+
+    # Load the state dictionary into the model
+    medsam_model.load_state_dict(checkpoint)  # Load the saved weights
+
     medsam_model = medsam_model.to(device)
     medsam_model.eval()