Repository: MouseLand/cellpose (language: python)

Issue:
Processing Large 3D image
Hi @carsen-stringer ,
In an existing [issue thread](https://github.com/MouseLand/cellpose/issues/178) I found that you recommended :  _**“if you're running out of memory due to the size of the tiff we don't currently have a workaround for that -- we recommend splitting the stack and stitching the masks afterwards.”**_

I’m wondering how to achieve an optimal stitching of the masks (aka Label images, in which all the pixels of an object have the same value). Indeed stitching of overlapping intensity values looks natural but my first guess is that it is much more complicated for Labels, in 3D! Detecting objects in 3D in the overlapping area, detecting partially/fully overlapping objects, deciding which one to keep … removing objects touching the edge ?
Am I overthinking it? 

With @lacan we thought it might be “easier” to fuse the flow images ? maybe in temporary files to overcome memory issues and run the other section of CellPose after on the fused image, which might not be so memory hungry?
Best,
Romain



Discussion / hints:
The way you propose, via blind rectilinear cropping of the input, would require some sophisticated overlap-based label stitching/merging, quite similar to how cell linear tracking is done. However, if you are able to threshold your cells _well enough_ to mask out cell clusters, you could actually partition your data via bounding boxes around those clusters. The cells that are cut off in one slice would be fully represented in another, so you could safely throw out any 'edge' labels and avoid the 'stitching' problem altogether (in the sense of stitching together cut-off cells - the resulting 'full' labels could simply be tiled together to form the full label matrix). There is probably even an 'optimal' partition of the data (fewest number of bounding boxes within some size range covering the cell clusters). 
The overlap is very easy to achieve via dask. you can basically make cellpose be run via dask (that splits the files, that's how I'm deconvolving my 2k by 2k images).
However, the labels numbers will mismatch. You could add   results = image2_masks + np.max(image1masks) and so on, but the edges would still be badly labelled. 
If you come up with an idea I'm very interested to apply it as I really need this for live cell imaging panels.
Yes, the edge labels are the issue, but the method I proposed earlier would solve this problem. If you are able to partition the images such that every cell appears in a non-cropped sub-image, then you can stitch back the panels and only include labels that are not on the edge. This does assume that cells are not duplicated - i.e. they appear only once in full view in a single sub-image, and can appear on the edge of an adjacent sub-image any number of times. This will be a function of cell density and the overlap of sub-images. How dense are your fields of cells? It would be useful to post an example. 
[dask-image's ndimage.label](https://github.com/dask/dask-image/blob/ec39605527f2e41012df1ec7e50634ee46469588/dask_image/ndmeasure/__init__.py#L263) has an example of stitching masks that might be helpful.  After it finishes labelling the blocks, it creates creates a connected components graph.  

If you are using dask, you'd need to use map_overlap with large enough depth to overcome cellpose edge effects at the chunk's nominal edge. 
check out this issue and the advice @qiagu and @Michael-Shannon

closing duplicate here https://github.com/MouseLand/cellpose/issues/171
Hey everyone, I've spent a bit of time thinking about this, probably less time than all of you though. I think the best option (which is pretty similar to @romainGuiet 's proposal I think) would be
1. the user gives as input a dask array (or cellpose makes a dask array from a folder of tiffs). 
2. Cellpose runs on slices through this array like `x[:, i, :]` and the rest of the array is not accessed (already flow computation is tiled on the GPU so no need to worry about GPU memory if a slice is big as long as it will fit in RAM). 
3. Each slice -> flow, cellprob gets sent to a zarr array for temporary storage. This means that you are always using overlapping tiles in an optimal way when computing the flows and cell probabilities.
4. Dynamics are run on overlapping chunks of the zarr array and stitched together. Maybe @chrisroat and @joaomamede can contribute some code for this second step. Any other way to run the dynamics seems really hard and would be difficult to run on the GPU. 

You can see the refactoring of the code section here splitting up flow computation and mask computation for this purpose: https://github.com/MouseLand/cellpose/blob/ff46c5f4ebd3e3c580be62765e38dc8108fcd041/cellpose/models.py#L526. In fact you can input a dask array to cellpose right now and it won't error, but it also won't be memory efficient because I have not implemented these intermediate steps to be stored on disk rather than in RAM.

What do people think about this approach?
Hi @carsen-stringer and everyone. I like this solution! I should be able to contribute by user testing - I have a bunch of big microscopy data that at the moment, I just crop and leave as is prior to cellpose. M
@carsen-stringer The proposed solution sounds great to me. I saw similar idea, tile iterator, work in other place. Look forward to seeing the idea getting implemented soon.
>the user gives as input a dask array (or cellpose makes a dask array from a folder of tiffs).

I definitely upvote the "cellpose makes a dask array from a folder of tiffs" (which would keep compatibility with our Cellpose wrapper for Fiji)

Thank you again for making this tool and being responsive to adding features

Cheers
Awesome, thanks everyone for the feedback, will start implementing in June. In the meantime feel free to contribute pull requests for any of the steps if you're interested in contributing!
Hi All!    Thanks @carsen-stringer for pointing me to this thread.

I have built a dask distributed wrapper which runs Cellpose.eval on blocks of a large array in parallel; find [the pertinent code here](https://github.com/GFleishman/cellpose/blob/359e480335d68caa04c7caa6ff66a089c767b63f/cellpose/distribute.py#L1).

The details of setting up a dask cluster are implemented in [my ClusterWrap library](https://github.com/GFleishman/ClusterWrap/blob/2e865557d76a3a81f73792ab67c1370d4b7b70b8/ClusterWrap/clusters.py#L1). Some modifications to this would probably be required to run it on your own cluster - though likely not much especially if you have experience with [dask-jobqueue](https://jobqueue.dask.org/en/latest/). It should work out of the box for workstations.

If you have a look at my WIP dask implementation you'll see I haven't yet worked out merging the overlap regions.

While I agree with @romainGuiet that fusing flow fields and probabilities is easier than fusing masks (just weighted averaging), the dynamics portion would then run on the whole stitched array(?) which in the very large image case may take quite some time. I also agree with @kevinjohncutler that if blocks could be defined around discrete cell clusters then the mask-stitching problem is much easier (or non-existent), but there are probably many biological situations where this is not possible - mine for example (whole brain zebrafish expansion microscopy).

So I am proceeding with the approach described by @chrisroat and I will post screenshots of my first rounds of testing to inspect the stitch boundaries. So far, I have noticed that within the overlap region the left-tile-segmentation and right-tile-segmentation are pretty consistent but not exactly the same. It may be that simply merging connected components across the boundary has a small percentage of false merges and failed merges. If so - then some kind of merging of the entire overlap region (something like [Joint Label Fusion](https://ieeexplore.ieee.org/document/6226425)) might be worth considering.
Sorry I took a bit to reply:

example 1:
Have a look at 
[Libraries.py](https://github.com/joaomamede/livepanelize/blob/main/LivePanelize/Libraries.py)
and
[Panelize_LiveCell.py](https://github.com/joaomamede/livepanelize/blob/main/LivePanelize/Panelize_LiveCell.py)

This is how I stitched my files (with 0% overlap).


Example 2:
When I want to cut and use overlaps generated by dask:

```
chunk_size=(frames.sizes['z'],frames.sizes['y']//ydivide,frames.sizes['x']//xdivide)

def depth_divide(xdivide=1,ydivide=1):
    if xdivide >1 and ydivide >1:
        depthdivide = (0,64,64)
    elif xdivide > 1 and ydivide == 1:
            depthdivide = (0,0,64)
    else:
            depthdivide = (0,0,0)

    return depthdivide

arr = da.from_array(frames[time]           
                                    , chunks=chunk_size)

res[:,i,:,:] = da.map_overlap(deconv,arr, depth  = depthdivide ,boundary='reflect',
                                    **kwargs,
                                             ).compute(num_workers=1)
```

Fix patch:
diff --git a/cellpose/contrib/__init__.py b/cellpose/contrib/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/cellpose/contrib/distributed_segmentation.py b/cellpose/contrib/distributed_segmentation.py
new file mode 100644
index 00000000..eb8212f7
--- /dev/null
+++ b/cellpose/contrib/distributed_segmentation.py
@@ -0,0 +1,260 @@
+"""Segments detected regions in a chunked dask array.
+
+Uses overlapping chunks during segmentation, and determines how to link segments
+between neighboring chunks by examining the overlapping border regions.
+
+Heavily based on dask_image.ndmeasure.label, which uses non-overlapping blocks
+with a structuring element that links segments at the chunk boundaries.
+"""
+import functools
+import logging
+import operator
+
+import dask
+import dask.array as da
+import numpy as np
+
+class DistSegError(Exception):
+    """Error in image segmentation."""
+
+try:
+    from dask_image.ndmeasure._utils import _label
+    from sklearn import metrics as sk_metrics
+except ModuleNotFoundError as e:
+    raise DistSegError("Install 'cellpose[distributed]' for distributed segmentation dependencies") from e
+
+
+logger = logging.getLogger(__name__)
+
+
+def segment(
+    image,
+    channels,
+    model_type,
+    diameter,
+    fast_mode=False,
+    use_anisotropy=True,
+    iou_depth=2,
+    iou_threshold=0.7,
+):
+    """Use cellpose to segment nuclei in fluorescence data.
+
+    Parameters
+    ----------
+    image : array of shape (z, y, x, channel)
+        Image used for detection of objects
+    channels : array of int with size 2
+        See cellpose
+    model_type : str
+        "cyto" or "nuclei"
+    diameter : tuple of size 3
+        Approximate diameter (in pixels) of a segmented region, i.e. cell width
+    fast_mode : bool
+        In fast mode, network averaging, tiling, and augmentation are turned off.
+    use_anisotropy : bool
+        If true, use anisotropy parameter of cellpose
+    iou_depth: dask depth parameter
+        Number of pixels of overlap to use in intersection-over-union calculation when
+        linking segments across neighboring, overlapping dask chunk regions.
+    iou_threshold: float
+        Minimum intersection-over-union in neighboring, overlapping dask chunk regions
+        to be considered the same segment.  The region for calculating IOU is given by the
+        iou_depth parameter.
+
+    Returns:
+        segments : array of int32 with same shape as input
+            Each segmented cell is assigned a number and all its pixels contain that value (0 is background)
+    """
+    assert image.ndim == 4, image.ndim
+    assert image.shape[-1] in {1, 2}, image.shape
+    assert diameter[1] == diameter[2], diameter
+
+    diameter_yx = diameter[1]
+    anisotropy = diameter[0] / diameter[1] if use_anisotropy else None
+
+    image = da.asarray(image)
+    image = image.rechunk({-1: -1})  # color channel is chunked together
+
+    depth = tuple(np.ceil(diameter).astype(np.int64))
+    boundary = "reflect"
+
+    # No chunking in channel direction
+    image = da.overlap.overlap(image, depth + (0,), boundary)
+
+    block_iter = zip(
+        np.ndindex(*image.numblocks),
+        map(
+            functools.partial(operator.getitem, image),
+            da.core.slices_from_chunks(image.chunks),
+        ),
+    )
+
+    labeled_blocks = np.empty(image.numblocks[:-1], dtype=object)
+    total = None
+    for index, input_block in block_iter:
+        labeled_block, n = dask.delayed(segment_chunk, nout=2)(
+            input_block,
+            channels,
+            model_type,
+            diameter_yx,
+            anisotropy,
+            fast_mode,
+            index,
+        )
+
+        shape = input_block.shape[:-1]
+        labeled_block = da.from_delayed(labeled_block, shape=shape, dtype=np.int32)
+
+        n = dask.delayed(np.int32)(n)
+        n = da.from_delayed(n, shape=(), dtype=np.int32)
+
+        total = n if total is None else total + n
+
+        block_label_offset = da.where(labeled_block > 0, total, np.int32(0))
+        labeled_block += block_label_offset
+
+        labeled_blocks[index[:-1]] = labeled_block
+        total += n
+
+    # Put all the blocks together
+    block_labeled = da.block(labeled_blocks.tolist())
+
+    depth = da.overlap.coerce_depth(len(depth), depth)
+
+    if np.prod(block_labeled.numblocks) > 1:
+        iou_depth = da.overlap.coerce_depth(len(depth), iou_depth)
+
+        if any(iou_depth[ax] > depth[ax] for ax in depth.keys()):
+            raise DistSegError("iou_depth (%s) > depth (%s)" % (iou_depth, depth))
+
+        trim_depth = {k: depth[k] - iou_depth[k] for k in depth.keys()}
+        block_labeled = da.overlap.trim_internal(
+            block_labeled, trim_depth, boundary=boundary
+        )
+        block_labeled = link_labels(
+            block_labeled,
+            total,
+            iou_depth,
+            iou_threshold=iou_threshold,
+        )
+
+        block_labeled = da.overlap.trim_internal(
+            block_labeled, iou_depth, boundary=boundary
+        )
+
+    else:
+        block_labeled = da.overlap.trim_internal(
+            block_labeled, depth, boundary=boundary
+        )
+
+    return block_labeled
+
+
+def segment_chunk(
+    chunk,
+    channels,
+    model_type,
+    diameter_yx,
+    anisotropy,
+    fast_mode,
+    index,
+):
+    """Perform segmentation on an individual chunk."""
+    # Cellpose seems to have some randomness, which is made deterministic by using the block
+    # details as a random seed.
+    np.random.seed(index)
+
+    from cellpose import models
+
+    model = models.Cellpose(gpu=True, model_type=model_type, net_avg=not fast_mode)
+
+    logger.info("Evaluating model")
+    segments, _, _, _ = model.eval(
+        chunk,
+        channels=channels,
+        z_axis=0,
+        channel_axis=3,
+        diameter=diameter_yx,
+        do_3D=True,
+        anisotropy=anisotropy,
+        net_avg=not fast_mode,
+        augment=not fast_mode,
+        tile=not fast_mode,
+    )
+    logger.info("Done segmenting chunk")
+
+    return segments.astype(np.int32), segments.max()
+
+
+def link_labels(block_labeled, total, depth, iou_threshold=1):
+    """
+    Build a label connectivity graph that groups labels across blocks,
+    use this graph to find connected components, and then relabel each
+    block according to those.
+    """
+    label_groups = label_adjacency_graph(block_labeled, total, depth, iou_threshold)
+    new_labeling = _label.connected_components_delayed(label_groups)
+    return _label.relabel_blocks(block_labeled, new_labeling)
+
+
+def label_adjacency_graph(labels, nlabels, depth, iou_threshold):
+    all_mappings = [da.empty((2, 0), dtype=np.int32, chunks=1)]
+
+    slices_and_axes = get_slices_and_axes(labels.chunks, labels.shape, depth)
+    for face_slice, axis in slices_and_axes:
+        face = labels[face_slice]
+        mapped = _across_block_iou_delayed(face, axis, iou_threshold)
+        all_mappings.append(mapped)
+
+    i, j = da.concatenate(all_mappings, axis=1)
+    result = _label._to_csr_matrix(i, j, nlabels + 1)
+    return result
+
+
+def _across_block_iou_delayed(face, axis, iou_threshold):
+    """Delayed version of :func:`_across_block_label_grouping`."""
+    _across_block_label_grouping_ = dask.delayed(_across_block_label_iou)
+    grouped = _across_block_label_grouping_(face, axis, iou_threshold)
+    return da.from_delayed(grouped, shape=(2, np.nan), dtype=np.int32)
+
+
+def _across_block_label_iou(face, axis, iou_threshold):
+    unique = np.unique(face)
+    face0, face1 = np.split(face, 2, axis)
+
+    intersection = sk_metrics.confusion_matrix(face0.reshape(-1), face1.reshape(-1))
+    sum0 = intersection.sum(axis=0, keepdims=True)
+    sum1 = intersection.sum(axis=1, keepdims=True)
+
+    # Note that sum0 and sum1 broadcast to square matrix size.
+    union = sum0 + sum1 - intersection
+
+    # Ignore errors with divide by zero, which the np.where sets to zero.
+    with np.errstate(divide="ignore", invalid="ignore"):
+        iou = np.where(intersection > 0, intersection / union, 0)
+
+    labels0, labels1 = np.nonzero(iou >= iou_threshold)
+
+    labels0_orig = unique[labels0]
+    labels1_orig = unique[labels1]
+    grouped = np.stack([labels0_orig, labels1_orig])
+
+    valid = np.all(grouped != 0, axis=0)  # Discard any mappings with bg pixels
+    return grouped[:, valid]
+
+
+def get_slices_and_axes(chunks, shape, depth):
+    ndim = len(shape)
+    depth = da.overlap.coerce_depth(ndim, depth)
+    slices = da.core.slices_from_chunks(chunks)
+    slices_and_axes = []
+    for ax in range(ndim):
+        for sl in slices:
+            if sl[ax].stop == shape[ax]:
+                continue
+            slice_to_append = list(sl)
+            slice_to_append[ax] = slice(
+                sl[ax].stop - 2 * depth[ax], sl[ax].stop + 2 * depth[ax]
+            )
+            slices_and_axes.append((tuple(slice_to_append), ax))
+    return slices_and_axes
diff --git a/setup.py b/setup.py
index 88ea2cc4..183f333e 100644
--- a/setup.py
+++ b/setup.py
@@ -21,6 +21,12 @@
         'sphinx_rtd_theme',
       ]
 
+distributed_deps = [
+        'dask',
+        'dask_image',
+        'scikit-learn',
+]
+
 try:
     import torch
     a = torch.ones(2, 3)
@@ -54,7 +60,8 @@
     ],
     extras_require = {
       'docs': docs_deps,
-      'gui': gui_deps
+      'gui': gui_deps,
+      'distributed': distributed_deps,
     },
     include_package_data=True,
     classifiers=(
diff --git a/tox.ini b/tox.ini
index 6dfef07a..28ccce44 100644
--- a/tox.ini
+++ b/tox.ini
@@ -24,9 +24,9 @@ passenv =
     DISPLAY XAUTHORITY
     NUMPY_EXPERIMENTAL_ARRAY_FUNCTION
     PYVISTA_OFF_SCREEN
-extras = gui
+extras = gui,distributed
 deps = 
-    .[gui]
+    .[gui,distributed]
     pytest  # https://docs.pytest.org/en/latest/contents.html
     pytest-cov  # https://pytest-cov.readthedocs.io/en/latest/
     pytest-xvfb ; sys_platform == 'linux'


Test patch:
diff --git a/tests/contrib/test_distributed_segmentation.py b/tests/contrib/test_distributed_segmentation.py
new file mode 100644
index 00000000..189cc7ce
--- /dev/null
+++ b/tests/contrib/test_distributed_segmentation.py
@@ -0,0 +1,190 @@
+"""Tests for distributed_segmentation.py module."""
+import pathlib
+
+import dask
+import dask.array as da
+import numpy as np
+import pytest
+import tifffile
+from dask.array.utils import assert_eq
+
+from cellpose.contrib import distributed_segmentation
+
+TESTDATA = pathlib.Path(__file__).parent / "testdata"
+
+def test_segment_nucleus():
+    """Tests segmentation on dask array."""
+    image = tifffile.imread(TESTDATA / "segment_80x224x448_input.tiff")
+    image = image[..., None]
+    image = da.from_array(image).rechunk({2: 224})
+
+    actual = distributed_segmentation.segment(
+        image,
+        [0, 0],
+        "cyto",
+        diameter=(10, 30, 30),
+        fast_mode=True,
+        use_anisotropy=False,
+    )
+
+    # Assume segmenting on a single GPU and process each chunk serially.
+    with dask.config.set(scheduler="synchronous"):
+        actual = actual.compute()
+
+    expected = tifffile.imread(TESTDATA / "segment_80x224x448_expected.tiff")
+    assert_eq(actual, expected)
+
+
+@pytest.mark.parametrize(
+    "data,expected",
+    [
+        ([0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0]),
+        ([0, 1, 1, 2, 2, 0], [0, 1, 1, 1, 1, 0]),
+        (
+            [
+                [0, 1, 1, 2, 2, 0],
+                [0, 1, 1, 2, 2, 0],
+            ],
+            [
+                [0, 1, 1, 1, 1, 0],
+                [0, 1, 1, 1, 1, 0],
+            ],
+        ),
+        (
+            [
+                [0, 1, 1, 2, 2, 0],
+                [0, 3, 3, 4, 4, 0],
+            ],
+            [
+                [0, 1, 1, 1, 1, 0],
+                [0, 2, 2, 2, 2, 0],
+            ],
+        ),
+        (
+            [
+                [0, 1, 0],
+                [0, 1, 0],
+                [0, 1, 0],
+                [0, 2, 0],
+                [0, 2, 0],
+                [0, 0, 0],
+            ],
+            [
+                [0, 1, 0],
+                [0, 1, 0],
+                [0, 1, 0],
+                [0, 1, 0],
+                [0, 1, 0],
+                [0, 0, 0],
+            ],
+        ),
+        (
+            [
+                [0, 1, 1, 2, 2, 0],
+                [0, 3, 3, 4, 4, 0],
+                [0, 3, 3, 4, 4, 0],
+                [0, 5, 5, 6, 6, 0],
+                [0, 5, 5, 6, 6, 0],
+                [0, 7, 7, 8, 8, 0],
+            ],
+            [
+                [0, 1, 1, 1, 1, 0],
+                [0, 2, 2, 2, 2, 0],
+                [0, 2, 2, 2, 2, 0],
+                [0, 2, 2, 2, 2, 0],
+                [0, 2, 2, 2, 2, 0],
+                [0, 3, 3, 3, 3, 0],
+            ],
+        ),
+        (
+            [
+                [1, 2, 3, 7, 8, 9],
+                [4, 5, 5, 10, 10, 11],
+                [6, 5, 5, 10, 10, 12],
+                [13, 14, 14, 19, 19, 20],
+                [15, 14, 14, 19, 19, 21],
+                [16, 17, 18, 22, 23, 24],
+            ],
+            [
+                [1, 2, 3, 2, 3, 7],
+                [4, 5, 5, 5, 5, 8],
+                [6, 5, 5, 5, 5, 9],
+                [4, 5, 5, 5, 5, 8],
+                [6, 5, 5, 5, 5, 9],
+                [10, 11, 12, 11, 12, 13],
+            ],
+        ),
+    ],
+)
+def test_link_labels(data, expected):
+    data = np.array(data, dtype=np.int32)
+    nlabels = np.max(data)
+    data = da.asarray(data, chunks=3)
+    result = distributed_segmentation.link_labels(data, nlabels, depth=1)
+    assert_eq(result, np.array(expected, dtype=np.int32))
+
+
+def test_link_labels_complex_case():
+    """Case with asymmetric depth.
+
+    Core is 4x6:
+    111110
+    111110
+    111022
+    011033
+    """
+    data = [
+        [0, 0, 0, 1, 1, 0, 0, 3, 3, 0, 0, 0, 0, 0],
+        [0, 0, 1, 1, 1, 1, 1, 3, 3, 3, 3, 0, 0, 0],
+        [0, 0, 1, 1, 1, 1, 1, 3, 3, 3, 3, 0, 0, 0],
+        [0, 0, 1, 1, 1, 0, 2, 3, 3, 0, 4, 4, 4, 4],
+        [0, 0, 5, 5, 5, 5, 5, 9, 9, 9, 9, 0, 0, 11],
+        [0, 0, 5, 5, 5, 0, 6, 9, 9, 0, 10, 10, 11, 11],
+        [0, 0, 0, 5, 5, 0, 7, 9, 9, 0, 11, 11, 11, 11],
+        [0, 0, 0, 8, 8, 0, 7, 12, 12, 0, 11, 13, 13, 14],
+    ]
+    data = np.array(data, dtype=np.int32)
+    data = da.asarray(data, chunks=(4, 7))
+
+    expected = [
+        [0, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 0, 0, 0],
+        [0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
+        [0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
+        [0, 0, 1, 1, 1, 0, 2, 1, 1, 0, 2, 2, 2, 2],
+        [0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 3],
+        [0, 0, 1, 1, 1, 0, 2, 1, 1, 0, 2, 2, 3, 3],
+        [0, 0, 0, 1, 1, 0, 3, 1, 1, 0, 3, 3, 3, 3],
+        [0, 0, 0, 4, 4, 0, 3, 4, 4, 0, 3, 5, 5, 6],
+    ]
+    result = distributed_segmentation.link_labels(data, np.max(data), depth=(1, 2))
+    assert_eq(result, np.array(expected, dtype=np.int32))
+
+
+def test_link_labels_threshold():
+    data = [
+        [0, 1, 1, 2, 2, 0],
+        [0, 1, 1, 2, 2, 0],
+        [0, 1, 1, 3, 3, 0],
+    ]
+    data = np.array(data, dtype=np.int32)
+    data = da.asarray(data, chunks=3)
+
+    expected = [
+        [0, 1, 1, 2, 2, 0],
+        [0, 1, 1, 2, 2, 0],
+        [0, 1, 1, 3, 3, 0],
+    ]
+    result = distributed_segmentation.link_labels(
+        data, np.max(data), depth=1, iou_threshold=0.8
+    )
+    assert_eq(result, np.array(expected, dtype=np.int32))
+
+    expected = [
+        [0, 1, 1, 1, 1, 0],
+        [0, 1, 1, 1, 1, 0],
+        [0, 1, 1, 2, 2, 0],
+    ]
+    result = distributed_segmentation.link_labels(
+        data, np.max(data), depth=1, iou_threshold=0.5
+    )
+    assert_eq(result, np.array(expected, dtype=np.int32))
diff --git a/tests/contrib/testdata/segment_80x224x448_expected.tiff b/tests/contrib/testdata/segment_80x224x448_expected.tiff
new file mode 100644
index 00000000..038849b3
Binary files /dev/null and b/tests/contrib/testdata/segment_80x224x448_expected.tiff differ
diff --git a/tests/contrib/testdata/segment_80x224x448_input.tiff b/tests/contrib/testdata/segment_80x224x448_input.tiff
new file mode 100644
index 00000000..a424f15b
Binary files /dev/null and b/tests/contrib/testdata/segment_80x224x448_input.tiff differ