Update cellemetry/services/sam.py

cellemetry/services/sam.py

@@ -1,5 +1,6 @@
 """
-SAM3 segmentation execution - Optimized for Speed.
+SAM3 segmentation execution - Optimized & Simplified.
+Removed morphological filtering for maximum recall and speed.
 """
 import matplotlib
 matplotlib.use('Agg')
@@ -10,14 +11,11 @@ import torchvision
 import numpy as np
 import time
 from PIL import Image
-from skimage.measure import regionprops
 
 from ..config.schemas import ComponentRequest
 from ..config.dependencies import AnalysisDeps
 
-MIN_SOLIDITY = 0.50
-MIN_CIRCULARITY = 0.1
-MAX_DIMENSION = 1024  # <-- SPEED OPTIMIZATION: Downscale large images
+MAX_DIMENSION = 1024  # Speed optimization: Downscale large images
 
 # Use /tmp for all outputs
 OUTPUT_DIR = "/tmp"
@@ -27,6 +25,8 @@ def execute_segmentation(deps: AnalysisDeps, request: ComponentRequest) -> str:
     Execute SAM3 segmentation for the given component request.
     """
     t_start = time.time()
+    
+    # SIMPLIFIED PROMPT: Color + Entity only (e.g., "green cell")
     text_prompt = f"{request.color} {request.entity}"
     print(f"\n[Engine] Segmenting: '{text_prompt}' ({len(request.bboxes)} boxes).")
 
@@ -36,7 +36,7 @@ def execute_segmentation(deps: AnalysisDeps, request: ComponentRequest) -> str:
     except Exception as e:
         return f"Error loading image: {e}"
 
-    # 2. SPEED FIX: Resize image if too large
+    # 2. Resize image if too large (Critical for speed)
     w, h = raw_image.size
     scale_factor = 1.0
     if max(w, h) > MAX_DIMENSION:
@@ -45,13 +45,11 @@ def execute_segmentation(deps: AnalysisDeps, request: ComponentRequest) -> str:
         new_h = int(h * scale_factor)
         raw_image = raw_image.resize((new_w, new_h), Image.Resampling.LANCZOS)
         print(f"[Engine] ⚡ Resized image from {w}x{h} to {new_w}x{new_h} (Speedup)")
-        # Update width/height for box calculations below
         w, h = new_w, new_h
 
     # 3. Convert normalized coords (0-1000) to pixel coords
     sam_input_boxes = []
     for box in request.bboxes:
-        # Scale coords to the (possibly resized) image dimensions
         y_min = (box.ymin / 1000) * h
         x_min = (box.xmin / 1000) * w
         y_max = (box.ymax / 1000) * h
@@ -83,7 +81,6 @@ def execute_segmentation(deps: AnalysisDeps, request: ComponentRequest) -> str:
         return_tensors="pt"
     ).to(deps.device)
 
-    # Use inference_mode for slight speedup over no_grad
     with torch.inference_mode():
         outputs = deps.sam_model(**inputs)
 
@@ -95,62 +92,23 @@ def execute_segmentation(deps: AnalysisDeps, request: ComponentRequest) -> str:
     
     print(f"[Engine] Inference took {time.time() - t_inf:.2f}s")
 
-    # 5. Morphology filtering (Optimized)
-    t_filter = time.time()
-    keep_indices_morph = []
-    
-    # Pre-fetch masks to cpu/numpy once
-    all_masks_np = results["masks"].detach().cpu().numpy().squeeze()
-    if all_masks_np.ndim == 2: # Handle single mask case
-        all_masks_np = all_masks_np[np.newaxis, ...]
+    # 5. REMOVED: Morphological Filtering (Solidity/Circularity)
+    # We now trust the SAM model and the Agent's bounding box completely.
 
-    for mask_np in all_masks_np:
-        mask_int = mask_np.astype(int)
-        
-        # Optimization: fast skip if mask is too small (noise)
-        if np.sum(mask_int) < 50:
-            keep_indices_morph.append(False)
-            continue
-
-        props = regionprops(mask_int)
-        if not props:
-            keep_indices_morph.append(False)
-            continue
-
-        prop = props[0]
-        # Fast calc circularity
-        perimeter = prop.perimeter
-        if perimeter == 0:
-            keep_indices_morph.append(False)
-            continue
-            
-        circularity = (4 * np.pi * prop.area) / (perimeter ** 2)
-        keep_indices_morph.append(prop.solidity > MIN_SOLIDITY and circularity > MIN_CIRCULARITY)
-
-    if any(keep_indices_morph):
-        keep_indices_tensor = torch.tensor(keep_indices_morph, device=results["masks"].device)
-        results = _filter_results(results, keep_indices_tensor)
-
-    print(f"[Engine] Filtering took {time.time() - t_filter:.2f}s")
-
-    # 6. NMS
+    # 6. NMS (Keep this to remove duplicate detections on the same object)
     pred_boxes = results["boxes"]
     pred_scores = results["scores"]
     if len(pred_scores) > 1:
         keep_indices_nms = torchvision.ops.nms(pred_boxes, pred_scores, iou_threshold=0.3)
         results = _filter_results(results, keep_indices_nms)
 
-    # 7. Save outputs (If resized, we must upscale masks back to original? 
-    # For demo purposes, we save the resized masks to keep things fast and aligned with the plot)
+    # 7. Save outputs
     _save_plot(raw_image, results, sam_input_boxes, text_prompt, plot_filename)
 
     mask_count = len(results['masks'])
     if mask_count > 0:
         masks_list = [m.cpu().numpy().squeeze() for m in results['masks']]
         masks_array = np.array(masks_list)
-        # If we resized, the stats (area) will be in resized pixels. 
-        # Ideally we'd resize masks back, but for a demo, just warn or accept.
-        # Alternatively, save the scale factor to adjust stats later.
         np.savez_compressed(data_filename, masks=masks_array)
     else:
         np.savez_compressed(data_filename, masks=np.array([]))
@@ -190,15 +148,12 @@ def _save_plot(image, results, boxes, label, filename):
         
         ax.imshow(composite)
 
-    # 2. Draw Bounding Boxes (NEW)
-    # 'boxes' comes in as [[xmin, ymin, xmax, ymax], ...]
+    # 2. Draw Bounding Boxes
     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='--'