From d24a6a3a0629607568c862bcb4a3ded432ef2540 Mon Sep 17 00:00:00 2001 From: Laurent Guerard Date: Tue, 30 Jun 2026 08:29:23 +0200 Subject: [PATCH] =?UTF-8?q?fix(core):=20=F0=9F=90=9B=20make=20per-tile=20l?= =?UTF-8?q?abels=20globally=20unique=20before=20merging?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Tiles each emit local labels (1..N), so every tile has a "1". The merge keyed objects by value, fusing unrelated objects from different tiles — multi-object, multi-tile segmentations collapsed (often to a single label). This is the package's core promise, so it must always hold. Add a streaming first pass in zarr_native_merge (_make_globally_unique) that renumbers each chunk's labels into a fresh contiguous global range before the boundary stitch. Unique AND compact: max_label == n_objects, so the relabel LUT stays O(objects), not O(voxels). One chunk in RAM at a time — no OOM (never more than segmentation already loads per tile). The merge output now matches the staged dtype so wide ids never truncate. Fixes both tile_process and the per-tile API (stage_tile). Regression test added: four separate objects across four tiles stay four labels. Co-Authored-By: Claude Opus 4.8 --- src/patchworks/_distributed.py | 5 ++- src/patchworks/_merge.py | 79 ++++++++++++++++++++++++++++------ tests/test_distributed.py | 30 +++++++++++++ 3 files changed, 100 insertions(+), 14 deletions(-) diff --git a/src/patchworks/_distributed.py b/src/patchworks/_distributed.py index 822a182..255433c 100644 --- a/src/patchworks/_distributed.py +++ b/src/patchworks/_distributed.py @@ -68,7 +68,8 @@ def create_stage( component : str, optional Array name inside the store (default ``"staged"``). dtype : data-type, optional - Label dtype (default ``int32``). + Label dtype (default ``int32``). Tiles write local labels; the merge's + first pass renumbers them to a compact global range that fits int32. Returns ------- @@ -135,6 +136,8 @@ def stage_tile( slice(left, out.shape[i] - right) for i, (left, right) in enumerate(trims) ) + # Local labels (1..N) are fine here — they collide across tiles, but the + # merge's first pass makes them globally unique before stitching. dst = zarr.open_group(str(stage_path), mode="r+")[component] dst[sl] = out[sel].astype(dst.dtype) return index diff --git a/src/patchworks/_merge.py b/src/patchworks/_merge.py index 109c1be..eb182fb 100644 --- a/src/patchworks/_merge.py +++ b/src/patchworks/_merge.py @@ -103,7 +103,7 @@ def _relabel_chunk_worker(chunk_slice: tuple) -> None: block = np.asarray(src[chunk_slice], dtype=np.int64) max_b = int(block.max()) if max_b == 0: - dst[chunk_slice] = block.astype(np.int32) + dst[chunk_slice] = block.astype(dst.dtype) return lut = _merge_lut if max_b < len(lut): @@ -111,7 +111,7 @@ def _relabel_chunk_worker(chunk_slice: tuple) -> None: else: ext = np.arange(len(lut), max_b + 1, dtype=np.int64) out = np.concatenate([lut, ext])[block] - dst[chunk_slice] = out.astype(np.int32) + dst[chunk_slice] = out.astype(dst.dtype) def _boundary_face_specs( @@ -238,9 +238,9 @@ def _build_relabel_lut(pairs: np.ndarray, max_label: int) -> np.ndarray: def _create_zarr_label_array( - group: zarr.Group, name: str, shape: tuple, chunks: tuple + group: zarr.Group, name: str, shape: tuple, chunks: tuple, dtype=np.int32 ) -> zarr.Array: - """Create (replacing any existing) an int32 label array in *group*. + """Create (replacing any existing) a label array in *group*. Parameters ---------- @@ -252,6 +252,9 @@ def _create_zarr_label_array( Array shape. chunks : tuple Chunk shape. + dtype : data-type, optional + Label dtype (default ``int32``). Matched to the staged store so the + merge keeps wide (block-encoded) ids intact before they are compacted. Returns ------- @@ -261,14 +264,56 @@ def _create_zarr_label_array( if name in group: del group[name] if _ZARR_V3: - return group.create_array( - name, shape=shape, chunks=chunks, dtype=np.int32 - ) + return group.create_array(name, shape=shape, chunks=chunks, dtype=dtype) return group.zeros( - name, shape=shape, chunks=chunks, dtype=np.int32, overwrite=True + name, shape=shape, chunks=chunks, dtype=dtype, overwrite=True ) +def _make_globally_unique(arr, shape: tuple, chunk_shape: tuple) -> int: + """Renumber per-chunk local labels to a globally unique, compact range. + + Each tile writes local labels (``1..N``) that repeat across tiles, so the + same value means different objects in different tiles. This streams the + chunks in row-major order and remaps every chunk's non-zero labels to a + fresh contiguous block ``[base+1, …]``, leaving the store globally unique + with ``max_label == total objects`` (compact, so the relabel LUT stays + ``O(n_objects)``). Background (0) stays 0. In place, one chunk in RAM. + + Parameters + ---------- + arr : zarr.Array + Writable staged label array. + shape : tuple + Array shape. + chunk_shape : tuple + Chunk (= tile) shape. + + Returns + ------- + int + New maximum label (number of objects across all tiles, before the + cross-boundary merge fuses touching pairs). + """ + n_per_dim = [(s + c - 1) // c for s, c in zip(shape, chunk_shape)] + base = 0 + for idx in _iproduct(*[range(n) for n in n_per_dim]): + sl = tuple( + slice(i * c, min((i + 1) * c, s)) + for i, c, s in zip(idx, chunk_shape, shape) + ) + block = np.asarray(arr[sl]) + uniq = np.unique(block) + uniq = uniq[uniq > 0] + if uniq.size == 0: + continue + lut = np.zeros(int(uniq[-1]) + 1, dtype=np.int64) + lut[uniq] = np.arange(1, uniq.size + 1) + base + arr[sl] = lut[block].astype(arr.dtype) + base += int(uniq.size) + return base + + def zarr_native_merge( staged_path: str, staged_component: str, @@ -302,13 +347,17 @@ def zarr_native_merge( ------- None """ - root = zarr.open_group(staged_path, mode="r") + root = zarr.open_group(staged_path, mode="r+") arr = root[staged_component] shape, chunk_shape = arr.shape, arr.chunks - max_label = int( - da.from_zarr(staged_path, component=staged_component).max().compute() - ) + # Pass 0 — make labels globally unique. Each tile writes local labels + # (1..N) that collide across tiles (every tile has a "1"); without this the + # boundary merge would fuse unrelated objects that merely share a value. + # Stream chunks in order, remapping each chunk's labels into a fresh + # contiguous range [base+1, …] — unique *and* compact (so the relabel LUT + # stays O(n_objects), not O(n_voxels)). + max_label = _make_globally_unique(arr, shape, chunk_shape) logger.info( "zarr_native_merge: shape=%s chunks=%s max_label=%d", shape, @@ -330,7 +379,11 @@ def zarr_native_merge( ) out_root = zarr.open_group(out_path, mode="a") - _create_zarr_label_array(out_root, out_component, shape, chunk_shape) + # Match the staged dtype so block-encoded (wide) ids survive the merge; + # they are compacted to a small range afterwards by relabel_sequential_zarr. + _create_zarr_label_array( + out_root, out_component, shape, chunk_shape, dtype=arr.dtype + ) n_per_dim = [(s + c - 1) // c for s, c in zip(shape, chunk_shape)] chunk_slices = [ diff --git a/tests/test_distributed.py b/tests/test_distributed.py index 0a0decf..34044d5 100644 --- a/tests/test_distributed.py +++ b/tests/test_distributed.py @@ -43,3 +43,33 @@ def test_stage_then_merge_stitches_boundary(tmp_path): ).compute() ids = np.unique(merged[merged > 0]) assert ids.size == 1, f"object split into {ids.size} labels" + + +def test_separate_objects_keep_distinct_labels(tmp_path): + """Different objects in different tiles must NOT collapse to one label. + + Each tile produces local labels (1..N); without the merge's global-uniqueness + pass every tile's "1" would fuse into a single object. + """ + img = np.zeros((16, 32), "uint16") + # four separate objects, one per (16x16) tile, none touching a boundary + img[3:6, 3:6] = 500 + img[3:6, 19:22] = 500 + img[10:13, 3:6] = 500 + img[10:13, 19:22] = 500 + + stage = str(tmp_path / "stage.zarr") + tile = (16, 16) + create_stage(stage, img.shape, tile) + for i in range(len(spatial_tiles(img.shape, tile))): + stage_tile(img, _fn, stage, i, tile_shape=tile, overlap=4) + + merged = merge_tile_labels( + stage, + write_to=str(tmp_path / "out.zarr"), + input_component="staged", + sequential_labels=True, + ).compute() + ids = np.unique(merged[merged > 0]) + assert ids.size == 4, f"expected 4 distinct objects, got {ids.size}" + assert set(ids.tolist()) == {1, 2, 3, 4} # contiguous after relabel