Update cellemetry/services/sam.py

cellemetry/services/sam.py

@@ -4,6 +4,7 @@ SAM3 segmentation execution - Optimized for Speed.
 import matplotlib
 matplotlib.use('Agg')
 import matplotlib.pyplot as plt
+import matplotlib.patches as patches
 import torch
 import torchvision
 import numpy as np
@@ -169,13 +170,12 @@ def _filter_results(results, keep_indices):
 
 
 def _save_plot(image, results, boxes, label, filename):
-    """Save visualization of segmentation results."""
+    """Save visualization of segmentation results with bounding boxes."""
     fig, ax = plt.subplots(figsize=(10, 10))
     ax.imshow(image)
     
-    # Batch visualization for speed
+    # 1. Draw Segmentation Masks
     if len(results['scores']) > 0:
-        # Create a single composite mask image for faster plotting than individual ax.imshow calls
         H, W = results['masks'][0].shape[-2:]
         composite = np.zeros((H, W, 4))
         
@@ -190,7 +190,24 @@ def _save_plot(image, results, boxes, label, filename):
         
         ax.imshow(composite)
 
-    ax.set_title(f"{label}")
+    # 2. Draw Bounding Boxes (NEW)
+    # 'boxes' comes in as [[xmin, ymin, xmax, ymax], ...]
+    for box in boxes:
+        x_min, y_min, x_max, y_max = box
+        width = x_max - x_min
+        height = y_max - y_min
+        
+        # Create a Rectangle patch
+        # red border, no fill, thick line
+        rect = patches.Rectangle(
+            (x_min, y_min), width, height, 
+            linewidth=2, edgecolor='red', facecolor='none', linestyle='--'
+        )
+        ax.add_patch(rect)
+
+    ax.set_title(f"{label} (with Agent BBoxes)")
     ax.axis('off')
-    fig.savefig(filename)
+    
+    # Save tightly to remove whitespace
+    fig.savefig(filename, bbox_inches='tight', pad_inches=0)
     plt.close(fig)