Use waterz.face_merge.slice_overlaps in branch_merge, update large decode yaml

branch_merge.py now imports slice_overlaps from waterz.face_merge
instead of duplicating the overlap computation. Both border stitching
(large_decode) and z-slice resolution (branch_merge) share the same
implementation.

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

Author: Donglai Wei

connectomics/decoding/decoders/branch_merge.py

@@ -102,54 +102,10 @@ def _compute_z_extents(seg: np.ndarray) -> Tuple[Dict[int, int], Dict[int, int]]
 
 
 # ---------------------------------------------------------------------------
-# Slice overlap computation
+# Slice overlap computation — delegates to waterz.face_merge.slice_overlaps
 # ---------------------------------------------------------------------------
 
-
-def _slice_overlaps(
-    s0: np.ndarray,
-    s1: np.ndarray,
-    z_aff: Optional[np.ndarray] = None,
-) -> np.ndarray:
-    """Overlap statistics between two consecutive 2D label maps.
-
-    Returns (N, K) float64 array where each row is::
-
-        [id0, id1, size0, size1, overlap, mean_z_affinity?]
-    """
-    fg = (s0 > 0) & (s1 > 0)
-    ncols = 6 if z_aff is not None else 5
-    if not fg.any():
-        return np.empty((0, ncols), dtype=np.float64)
-
-    a = s0[fg].astype(np.int64)
-    b = s1[fg].astype(np.int64)
-
-    u0, c0 = np.unique(s0[s0 > 0], return_counts=True)
-    u1, c1 = np.unique(s1[s1 > 0], return_counts=True)
-    size0_map = dict(zip(u0.tolist(), c0.tolist()))
-    size1_map = dict(zip(u1.tolist(), c1.tolist()))
-
-    pairs = np.stack([a, b], axis=1)
-    unique_pairs, inverse, counts = np.unique(
-        pairs, axis=0, return_inverse=True, return_counts=True,
-    )
-
-    n = len(unique_pairs)
-    result = np.zeros((n, ncols), dtype=np.float64)
-    result[:, 0] = unique_pairs[:, 0]
-    result[:, 1] = unique_pairs[:, 1]
-    result[:, 2] = np.array([size0_map[int(i)] for i in unique_pairs[:, 0]])
-    result[:, 3] = np.array([size1_map[int(i)] for i in unique_pairs[:, 1]])
-    result[:, 4] = counts
-
-    if z_aff is not None:
-        aff_vals = z_aff[fg].astype(np.float64)
-        aff_sums = np.zeros(n, dtype=np.float64)
-        np.add.at(aff_sums, inverse, aff_vals)
-        result[:, 5] = aff_sums / counts
-
-    return result
+from waterz.face_merge import slice_overlaps as _slice_overlaps
 
 
 # ---------------------------------------------------------------------------

tutorials/waterz_decoding_large.yaml

@@ -23,28 +23,29 @@ description: >
 # ── Chunked decoding config ──────────────────────────────────────────────────
 large_decode:
   # Input affinity file (HDF5, dataset "main", shape (C,Z,Y,X))
-  affinity_path: ""       # SET THIS
+  affinity_path: "/projects/weilab/dataset/liconn/mansour/DL288B_251222S_cond5_40x_12tiles_round1_fused_488_crop512x1024x1024_ffn_sharp_tta_x8_prediction_uint8.h5"       # SET THIS
 
   # Workflow directory (shared filesystem for parallel workers)
-  workflow_root: ""       # SET THIS: e.g. /scratch/waterz_workflow/
+  workflow_root: "/projects/weilab/weidf/lib/pytorch_connectomics/outputs/neuron_liconn_mit/DL228B/"       # SET THIS: e.g. /scratch/waterz_workflow/
 
   # Chunk size in voxels (Z, Y, X). Each chunk is decoded independently.
   # Smaller chunks = less memory per worker but more border stitching.
   # Rule of thumb: ~256^3 for 16 GB RAM, ~512^3 for 64 GB RAM.
-  chunk_shape: [256, 512, 512]
+  chunk_shape: [80, 2066, 2066]
 
   # Waterz agglomeration parameters (same as decode_waterz kwargs)
   thresholds: [0.5]
   merge_function: aff85_his256
-  aff_threshold_low: 0.1
+  aff_threshold_low: 0.001
   aff_threshold_high: 0.999
   channel_order: xyz
 
-  # Border stitching parameters
-  border_min_overlap: 1          # min overlap pixels to consider a pair
-  border_one_sided_threshold: 0.9  # IOU for small-into-large merge at borders
-  border_iou_threshold: 0.0     # full Jaccard IOU threshold (0=disabled)
-  border_affinity_threshold: 0.0 # min boundary affinity (0=disabled)
+  # Border stitching parameters (same algorithm as branch_merge)
+  border_min_overlap: 10          # min overlap pixels to consider a pair
+  border_iou_threshold: 0.0       # full Jaccard IOU threshold (0=disabled)
+  border_one_sided_threshold: 0.9 # overlap/min_size for small-into-large merge
+  border_one_sided_min_size: 0    # min segment size in face for one-sided merge
+  border_affinity_threshold: 0.0  # min boundary affinity (0=disabled)
 
   # Output
   write_output: true             # assemble final volume