Drop branch_ prefix from merge params: unified naming with face_merge_pairs

branch_iou_threshold → iou_threshold, branch_best_buddy → best_buddy,
branch_one_sided_threshold → one_sided_threshold, etc. Same names as
face_merge_pairs and waterz_decoding_large.yaml. branch_merge remains
as the enable flag.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: Donglai Wei

connectomics/decoding/decoders/waterz.py

@@ -46,11 +46,11 @@ def decode_waterz(
     dust_merge_affinity: float = 0.0,
     dust_remove_size: int = 0,
     branch_merge: bool = False,
-    branch_iou_threshold: float = 0.5,
-    branch_best_buddy: bool = True,
-    branch_one_sided_threshold: float = 0.8,
-    branch_one_sided_min_size: int = 100,
-    branch_affinity_threshold: float = 0.0,
+    iou_threshold: float = 0.5,
+    best_buddy: bool = True,
+    one_sided_threshold: float = 0.8,
+    one_sided_min_size: int = 100,
+    affinity_threshold: float = 0.0,
     return_all_thresholds: bool = False,
     **kwargs: Any,
 ) -> "np.ndarray | Dict[float, np.ndarray]":
@@ -123,18 +123,16 @@ def decode_waterz(
         branch_merge: Enable branch merge postprocessing.  Resolves false
             splits by analyzing segment continuity across z-slices using
             IOU overlap, best-buddy matching, and one-sided IOU.
-            Inspired by em_pipeline branch resolution.  Default: False
-        branch_iou_threshold: Stage 1 threshold for full Jaccard IOU
-            merge between consecutive z-slices.  Default: 0.5
-        branch_best_buddy: Enable Stage 2 mutual best-match merge.
-            Default: True
-        branch_one_sided_threshold: Stage 3 one-sided IOU threshold.
-            Merge if ``overlap / min(size0, size1)`` exceeds this.
-            Set to 0 to disable.  Default: 0.8
-        branch_one_sided_min_size: Stage 3 minimum segment size in the
-            slice to be considered for one-sided merge.  Default: 100
-        branch_affinity_threshold: Minimum mean z-boundary affinity for
-            a branch merge to be accepted.  0 to disable.  Default: 0.0
+            Default: False
+        iou_threshold: Full Jaccard IOU threshold for branch merge.
+            Default: 0.5
+        best_buddy: Enable mutual best-match merge.  Default: True
+        one_sided_threshold: One-sided IOU threshold
+            (``overlap / min(size0, size1)``).  0 to disable.  Default: 0.8
+        one_sided_min_size: Minimum segment size in slice for one-sided
+            merge.  Default: 100
+        affinity_threshold: Minimum mean z-boundary affinity for a merge.
+            0 to disable.  Default: 0.0
         return_all_thresholds: If True and multiple thresholds are given,
             return a dict mapping each threshold to its segmentation.
             Otherwise return only the last threshold's result. Default: False
@@ -302,11 +300,11 @@ def decode_waterz(
             seg = _branch_merge(
                 seg,
                 affinities=affs,
-                iou_threshold=branch_iou_threshold,
-                best_buddy=branch_best_buddy,
-                one_sided_threshold=branch_one_sided_threshold,
-                one_sided_min_size=branch_one_sided_min_size,
-                affinity_threshold=branch_affinity_threshold,
+                iou_threshold=iou_threshold,
+                best_buddy=best_buddy,
+                one_sided_threshold=one_sided_threshold,
+                one_sided_min_size=one_sided_min_size,
+                affinity_threshold=affinity_threshold,
                 channel_order="zyx",  # already converted above
             )
             n_after = len(np.unique(seg)) - (1 if 0 in seg else 0)

connectomics/training/lightning/model.py

@@ -912,10 +912,10 @@ def _log_decode_experiment(
         if decode_params.get("branch_merge"):
             param_keys += [
                 "branch_merge",
-                "branch_iou_threshold",
-                "branch_best_buddy",
-                "branch_one_sided_threshold",
-                "branch_one_sided_min_size",
+                "iou_threshold",
+                "best_buddy",
+                "one_sided_threshold",
+                "one_sided_min_size",
             ]
         metric_keys = [
             "adapted_rand_error",

connectomics/training/lightning/utils.py

@@ -474,11 +474,11 @@ def _sanitize_decode_component(text: str) -> str:
 
     gated_value_groups = {
         "branch_merge": [
-            "branch_iou_threshold",
-            "branch_best_buddy",
-            "branch_one_sided_threshold",
-            "branch_one_sided_min_size",
-            "branch_affinity_threshold",
+            "iou_threshold",
+            "best_buddy",
+            "one_sided_threshold",
+            "one_sided_min_size",
+            "affinity_threshold",
         ],
         "dust_merge": [
             "dust_merge_size",

tutorials/neuron_snemi.yaml

@@ -92,12 +92,12 @@ test:
         dust_merge_size: 800
         dust_merge_affinity: 0.3
         dust_remove_size: 600
-        branch_merge: false              # enable branch merge
-        branch_iou_threshold: 0.5       # S1: full Jaccard IOU threshold
-        branch_best_buddy: true         # S2: mutual best-match merge
-        branch_one_sided_threshold: 0.8 # S3: one-sided IOU threshold (0 to disable)
-        branch_one_sided_min_size: 100  # S3: min segment size for one-sided merge
-        branch_affinity_threshold: 0.0  # min z-boundary affinity (0=no check)
+        branch_merge: false
+        iou_threshold: 0.5
+        best_buddy: true
+        one_sided_threshold: 0.8
+        one_sided_min_size: 100
+        affinity_threshold: 0.0
 
 
 # ============================================================================

tutorials/neuron_snemi_sdt_multitask.yaml

@@ -160,17 +160,17 @@ test:
           #use_seg_uint32: true      # uint32 segment IDs (2x less seg memory)
           thresholds: 0.4
           merge_function: aff85_his256
-          aff_threshold: [0.1, 0.999]
+          aff_threshold: [0.1, 0.9]
           border_threshold: 0
           dust_merge: true
           dust_merge_size: 800
           dust_merge_affinity: 0.3
           dust_remove_size: 600
           branch_merge: false
-          branch_iou_threshold: 0.5
-          branch_best_buddy: true
-          branch_one_sided_threshold: 0.8
-          branch_one_sided_min_size: 100
+          iou_threshold: 0.5
+          best_buddy: true
+          one_sided_threshold: 0.8
+          one_sided_min_size: 100
 
 # ============================================================================
 # Parameter tuning for waterz agglomeration thresholds (--mode tune)