Update app.py

app.py

@@ -60,6 +60,7 @@ def visualize(image, mask, box):
     buf.seek(0)
     return buf
 
+# Main function for Gradio app
 # Main function for Gradio app
 def process_images(file, x_min, y_min, x_max, y_max):
     image, H, W = load_image(file)
@@ -67,24 +68,20 @@ def process_images(file, x_min, y_min, x_max, y_max):
     # 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':
-        checkpoint = torch.load(model_checkpoint_path, map_location=torch.device('cpu'))
-    else:
-        checkpoint = torch.load(model_checkpoint_path)
+    # Define the checkpoint path
+    model_checkpoint_path = "medsam_vit_b.pth"  # Replace with the correct path to your checkpoint
 
-    # Create model instance
-    medsam_model = sam_model_registry['vit_b']  # Create model instance without checkpoint
+    # Create the model instance and load the checkpoint
+    medsam_model = sam_model_registry['vit_b'](checkpoint=model_checkpoint_path)
 
-    # Load the state dictionary into the model
-    medsam_model.load_state_dict(checkpoint)  # Load the saved weights
+    # If running on CPU, map the model to CPU
+    if device == 'cpu':
+        medsam_model = medsam_model.to(torch.device('cpu'))
 
-    medsam_model = medsam_model.to(device)
     medsam_model.eval()
 
     box = [x_min, y_min, x_max, y_max]
-    mask = medsam_inference(medsam_model, image, box, H, W, H, device) # Assuming target size is the same as the image height
+    mask = medsam_inference(medsam_model, image, box, H, W, H, device)  # Pass device to the inference function
 
     visualization = visualize(image, mask, box)
     return visualization.getvalue() # Returning the byte stream