fix autoboundary

S2FApp/ui/heatmaps.py

@@ -98,22 +98,21 @@ def make_annotated_heatmap_multi_regions(heatmap_rgb, masks, labels, cell_mask=N
 
 
 def add_cell_contour_to_fig(fig_pl, cell_mask, row=1, col=2):
-    """Add red contour overlay to Plotly heatmap subplot."""
+    """Add red contour overlay to Plotly heatmap subplot. Draws all contours (handles multiple disconnected regions)."""
     if cell_mask is None or not np.any(cell_mask > 0):
         return
     contours, _ = cv2.findContours(cell_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
     if not contours:
         return
-    # Use largest contour
-    cnt = max(contours, key=cv2.contourArea)
-    pts = cnt.squeeze()
-    if pts.ndim == 1:
-        pts = pts.reshape(1, 2)
-    x, y = pts[:, 0].tolist(), pts[:, 1].tolist()
-    if x[0] != x[-1] or y[0] != y[-1]:
-        x.append(x[0])
-        y.append(y[0])
-    fig_pl.add_trace(
-        go.Scatter(x=x, y=y, mode="lines", line=dict(color="red", width=4), showlegend=False),
-        row=row, col=col
-    )
+    for cnt in contours:
+        pts = cnt.squeeze()
+        if pts.ndim == 1:
+            pts = pts.reshape(1, 2)
+        x, y = pts[:, 0].tolist(), pts[:, 1].tolist()
+        if x[0] != x[-1] or y[0] != y[-1]:
+            x.append(x[0])
+            y.append(y[0])
+        fig_pl.add_trace(
+            go.Scatter(x=x, y=y, mode="lines", line=dict(color="red", width=4), showlegend=False),
+            row=row, col=col
+        )

S2FApp/utils/segmentation.py

@@ -7,10 +7,13 @@ from skimage.measure import label, regionprops
 
 
 def estimate_cell_mask(heatmap, sigma=2, min_size=200, exclude_full_image=True,
-                       threshold_relax=0.85, dilate_radius=4):
+                       threshold_relax=0.85, dilate_radius=4, min_area_ratio=0.2):
     """
     Estimate cell region from force map using Otsu thresholding and morphological cleanup.
 
+    Supports multiple disconnected regions (e.g., two cells): components whose area is
+    at least min_area_ratio of the largest are merged into the final mask.
+
     Args:
         heatmap: 2D float array [0, 1] - predicted force map
         sigma: Gaussian smoothing sigma to reduce noise. Default 2.
@@ -21,6 +24,9 @@ def estimate_cell_mask(heatmap, sigma=2, min_size=200, exclude_full_image=True,
             Default 0.85.
         dilate_radius: Radius to dilate mask outward to include surrounding pixels.
             Default 4.
+        min_area_ratio: Include components with area >= this fraction of the largest
+            component (0–1). E.g. 0.2 = include regions at least 20% the size of the
+            largest. Handles multiple disconnected force regions. Default 0.2.
 
     Returns:
         mask: Binary uint8 array, 1 = estimated cell, 0 = background
@@ -39,9 +45,10 @@ def estimate_cell_mask(heatmap, sigma=2, min_size=200, exclude_full_image=True,
     # Morphological cleanup
     mask = closing(mask, disk(5)).astype(np.uint8)
     mask = opening(mask, disk(3)).astype(np.uint8)
-    mask = remove_small_objects(mask.astype(bool), max_size=min_size - 1).astype(np.uint8)
+    mask = remove_small_objects(mask.astype(bool), min_size=min_size).astype(np.uint8)
 
-    # Select component: second largest if largest is whole image
+    # Select component(s): optionally exclude full-image background, then merge
+    # all significant components (handles multiple disconnected force regions)
     labeled = label(mask)
     props = list(regionprops(labeled))
 
@@ -51,12 +58,18 @@ def estimate_cell_mask(heatmap, sigma=2, min_size=200, exclude_full_image=True,
     props_sorted = sorted(props, key=lambda x: x.area, reverse=True)
     total_px = heatmap.shape[0] * heatmap.shape[1]
 
+    # Skip largest if it covers most of image (likely background)
     if exclude_full_image and len(props_sorted) >= 2 and props_sorted[0].area > 0.7 * total_px:
-        region = props_sorted[1]
-    else:
-        region = props_sorted[0]
+        props_sorted = props_sorted[1:]
+
+    # Reference area for "significant" components
+    ref_area = props_sorted[0].area
+    # Include all components with area >= min_area_ratio * ref_area
+    labels_to_keep = [p.label for p in props_sorted if p.area >= min_area_ratio * ref_area]
 
-    mask = (labeled == region.label).astype(np.uint8)
+    mask = np.zeros_like(labeled, dtype=np.uint8)
+    for lab in labels_to_keep:
+        mask[labeled == lab] = 1
 
     # Dilate to include surrounding pixels
     if dilate_radius > 0: