fixed scaling

app.py

@@ -74,30 +74,26 @@ def process_images(file, x_min, y_min, x_max, y_max):
     image, H, W = load_image(file)
     image_resized = transform.resize(image, (1024, 1024), anti_aliasing=True)
     image_resized = (image_resized - image_resized.min()) / np.clip(image_resized.max() - image_resized.min(), a_min=1e-8, a_max=None)
-    
-    # Check if CUDA is available, and set the device accordingly
-    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    image_tensor = torch.tensor(image_resized).float().permute(2, 0, 1).unsqueeze(0).to(device)
 
-    # Define the checkpoint path
+    # Initialize the MedSAM model and set the device
     model_checkpoint_path = "medsam_vit_b.pth"  # Replace with the correct path to your checkpoint
-
-    # Create the model instance and load the checkpoint
     medsam_model = sam_model_registry['vit_b'](checkpoint=model_checkpoint_path)
     medsam_model = medsam_model.to(device)
     medsam_model.eval()
 
-    # Convert image to tensor and move to the correct device
-    image_tensor = torch.tensor(image_resized).float().permute(2, 0, 1).unsqueeze(0).to(device)
-
     # Generate image embedding
     with torch.no_grad():
         img_embed = medsam_model.image_encoder(image_tensor)
 
-    # Calculate resized box coordinates and perform inference
+    # Calculate resized box coordinates
     scale_factors = np.array([1024 / W, 1024 / H, 1024 / W, 1024 / H])
     box_1024 = np.array([x_min, y_min, x_max, y_max]) * scale_factors
+
+    # Perform inference
     mask = medsam_inference(medsam_model, img_embed, box_1024, H, W)
 
+    # Visualization
     visualization = visualize(image, mask, [x_min, y_min, x_max, y_max])
     return visualization.getvalue()