<jupyter_start><jupyter_text>Distributed WatershedImplementation of a distributed watershed function using Dask. Implementation inspired by <NAME> [here](https://github.com/dask/dask-image/pull/99). The overall method relies on a two-pass watershed model. The first pass watershed is used to generate and share information about markers across chunk boundaries. The second pass then propagates that information.This implementation differs from v1. Whereas v1 shares marker information about the first pass watershed directly, this shares information about through a specially-labeled chunk-boundary basin.  Base Imports<jupyter_code>import dask
import dask.array as da
import numpy as np<jupyter_output><empty_output><jupyter_text>Helper Functions<jupyter_code>def compute_mem_mb(shape):
    '''Determines memory consumption of an array with shape in MB'''
    from functools import reduce
    from operator import mul
    count = reduce(mul, shape, 1)
    return count * 8 / ( 1024 ** 2 )

def display(image):
    '''Shows an image in the Jupyter notebook.'''
    from skimage import io
    io.imshow(image)
    io.show()

def frac_to_zscore(frac):
    from scipy.stats import norm
    return norm.ppf(frac)

def mask_overlap(chunk, depth, label):
    overlap = np.zeros_like(chunk)
    for i in range(0, chunk.ndim):
        idx = [slice(None)]*chunk.ndim
        r = np.array(range(0, chunk.shape[i]))
        idx[i] = (r[:depth], r[-depth:])
        overlap[tuple(idx)] = label
    return overlap

def build_full_markers(labels, depth, label):
    markers = mask_overlap(labels, depth, label)
    markers[labels > 0] = labels[labels > 0]
    return markers

def create_random_salt_image(fraction_salt, size):
    zscore = frac_to_zscore(1 - fraction_salt)
    salt = np.random.normal(0.0, 1.0, size) > zscore
    return salt

def remove_label(chunk, label):
    out = chunk.copy()
    out[out == label] = 0
    return out

def compose(fp_chunk, sp_chunk, mask_chunk):
    out = fp_chunk.copy()
    out[mask_chunk] = sp_chunk[mask_chunk]
    return out<jupyter_output><empty_output><jupyter_text>Image Geometry Definitions<jupyter_code>ndim = 2
size_len = 200
size = ndim * [size_len]
mem_mb = compute_mem_mb(size)
print("2D array size (MB): {:.2f}".format(mem_mb))
chunk_len = 400
chunks = ndim * [chunk_len]
chunk_mem_mb = compute_mem_mb(chunks)
print("Chunk size (MB): {:.2f}".format(chunk_mem_mb))<jupyter_output><empty_output><jupyter_text>Random Seed<jupyter_code>seed = 1
np.random.seed(seed=seed)<jupyter_output><empty_output><jupyter_text>Create Random "Salting" Image<jupyter_code>fraction_salt = 1e-2
salt = create_random_salt_image(fraction_salt, size)
#salt = np.flip(salt, axis=0)
display(salt)<jupyter_output><empty_output><jupyter_text>Determine EDT<jupyter_code>from scipy.ndimage.morphology import distance_transform_edt
edt = distance_transform_edt(~salt)
max_edt = edt.max()
edt = edt / max_edt
display(edt)<jupyter_output><empty_output><jupyter_text>Filter EDT Using H-Max/H-Dome<jupyter_code>from skimage.morphology import reconstruction
h = 1 / max_edt
h_seed = edt - h
hmax = reconstruction(h_seed, edt, method='dilation')
display(hmax)<jupyter_output><empty_output><jupyter_text>Prepare Marker Image<jupyter_code>from skimage.measure import label
ws_markers = label(salt)<jupyter_output><empty_output><jupyter_text>Watershed Transform<jupyter_code>from skimage.morphology import watershed
ws = watershed(hmax, markers=ws_markers)
display(ws)<jupyter_output><empty_output><jupyter_text>Prepare Dask Client<jupyter_code>from dask.distributed import Client
c = Client()
port = c.scheduler_info()['services']['dashboard']
print("Type `http://localhost:{port}` into the URL bar of your favorite browser to watch the following code in action on your machine in real time.".format(port=port))<jupyter_output><empty_output><jupyter_text>Create Dask ArraysHere we assume future users have access to distributed versions of h-max, EDT, and connected component labeling.<jupyter_code>depth = 1
hmax_da = da.from_array(hmax, chunks=chunks)
hmax_op = da.overlap.overlap(hmax_da, depth=depth, boundary='nearest')

boundary_label = ws_markers.max() + 1
ws_markers_da = da.from_array(ws_markers, chunks=chunks)
ws_markers_op = da.overlap.overlap(ws_markers_da, depth=depth, boundary=boundary_label)
ws_markers_op = ws_markers_op.map_blocks(lambda x: build_full_markers(x, depth, boundary_label), dtype=ws_markers_op.dtype)
display(ws_markers_op)

# fp = first pass
ws_fp = hmax_op.map_blocks(lambda x, y: watershed(x, markers=y), ws_markers_op, dtype=hmax_op.dtype)
display(ws_fp)<jupyter_output><empty_output><jupyter_text>Propagate Boundary Basin<jupyter_code>ws_markers_sp = ws_fp.map_blocks(lambda x: remove_label(x, boundary_label))
ws_markers_sp = da.overlap.trim_overlap(ws_markers_sp, depth=depth)
ws_markers_sp = da.overlap.overlap(ws_markers_sp, depth=depth, boundary='nearest')
display(ws_markers_sp)

ws_mask = ws_fp == boundary_label
display(ws_mask)

ws_sp = hmax_op.map_blocks(lambda x, y, z: watershed(x, markers=y, mask=z), ws_markers_sp, ws_mask, dtype=ws_fp.dtype)
display(ws_sp)<jupyter_output><empty_output><jupyter_text>Compose First and Second Passes Using Mask<jupyter_code>ws_final = ws_fp.map_blocks(lambda x, y, z: compose(x, y, z), ws_sp, ws_mask, dtype=ws_fp.dtype)
ws_final = da.overlap.trim_overlap(ws_final, depth=depth)
display(ws_final)<jupyter_output><empty_output><jupyter_text>ValidationNote there are some different basin assignments between the methods. Pay particular attention to the error-free strip along the bottom 1/4 of the image. That strip seems to stay through seed changes, and even flipping the salt image along the vertical axis.<jupyter_code>error = ~(ws == ws_final).compute()
display(error)

error_count = error.sum()
print("Error count: {:d}".format(error_count))
print("Error fraction: {:.3%}".format(error_count / error.size))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10782990212622343, 'fraction_numerical': 0.006918663516706041, 'mean_word_length': 6.40875, 'pattern_counts': {'":': 0, '<': 58, '<?xml version=': 0, '>': 61, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '14', 'dir': 'jupyter-structured-clean-dedup', 'id': '411016', 'license': 'Unlicense', 'n_tokens_mistral': 2140, 'n_tokens_neox': 2044, 'n_words': 478, 'path': 'dask_image_watershed_v2.ipynb', 'repo_name': 'wwarriner/dask_test', 'size': '5939'}