<jupyter_start><jupyter_text>3D Material Bead Packing XCT Simulation Instance Segmentation - Distributed Computation<jupyter_code>import os
import time

import itk
import numpy as np

from matplotlib import pyplot as plt
from itkwidgets import view, label_statistics
import itkwidgets

from dask.distributed import Client, LocalCluster
import dask.array as da
import dask
import dask_image.ndmeasure
import zarr
# Option 1: local cluster

# local_cluster = LocalCluster(n_workers=8, processes=False, memory_limit='4G')
# client = Client(local_cluster)
# client
# Option 2: NERSC Cori Slurm dask-mpi cluster
#
# See SC20_pyHPC/nersc/README.md

scheduler_file = os.path.join(os.environ["SCRATCH"], "scheduler.json")
dask.config.config["distributed"]["dashboard"]["link"] = "{JUPYTERHUB_SERVICE_PREFIX}proxy/{host}:{port}/status"

client = Client(scheduler_file=scheduler_file)
# client
file_name = '../data/bead_pack.tif'
beads = itk.imread(file_name)
beads_da = da.from_array(np.asarray(beads))
beads_da
slice_chunk_size = 100
beads_da = beads_da.rechunk((100,200,200))
beads_da
zarr_path = '../data/bead_pack.zarr'

beads_da.to_zarr(zarr_path, overwrite=True, compute=True)
beads = da.from_zarr(zarr_path)
beads
def processing_pipeline(image_chunk):
    import itk
    
    sharpened = itk.unsharp_mask_image_filter(image_chunk,
                                              sigma=2.0,
                                              threshold=10.0,
                                              amount=5.0)
    invert = itk.invert_intensity_image_filter(sharpened)
    thresholded = itk.otsu_threshold_image_filter(invert)
    
    StructuringElementType = itk.FlatStructuringElement[3]

    structuring_element = StructuringElementType.Ball(3, True)
    opened = itk.binary_morphological_opening_image_filter(thresholded,
                                                           kernel=structuring_element)
    filled = itk.binary_fillhole_image_filter(opened, fully_connected=True)
    
    distance = itk.signed_maurer_distance_map_image_filter(filled, inside_is_positive=False)
    
    watershed = itk.morphological_watershed_image_filter(distance, level=7.0)
    filled_cast = filled.astype(np.int16)
    watershed_mask = itk.mask_image_filter(watershed, filled_cast)

    return watershed_mask
start = time.time()

segmented = da.map_blocks(processing_pipeline,
              beads,
              dtype=np.int16)
segmented = segmented.compute()

elapsed = time.time() - start
print(elapsed, 'seconds')
view(beads, label_image=segmented, ui_collapsed=True)
start = time.time()

segmented = da.map_blocks(processing_pipeline,
              beads,
              dtype=np.int16)
segmented = segmented.compute()

elapsed = time.time() - start
print(elapsed, 'seconds')
chunks = list(beads.chunks)
slice_chunks = list(chunks[0])
chunk_shift = int(slice_chunk_size / 2)
original_chunk = slice_chunks[0]
slice_chunks[0] = chunk_shift
slice_chunks.append(original_chunk - chunk_shift)
chunks[0] = tuple(slice_chunks)
chunks
beads_rechunk = beads.rechunk(chunks)
beads_rechunk
start = time.time()

segmented_rechunk = da.map_blocks(processing_pipeline,
              beads_rechunk,
              dtype=np.int16)
segmented_rechunk = segmented_rechunk.compute()

elapsed = time.time() - start
print(elapsed, 'seconds')
view(beads_rechunk, label_image=segmented_rechunk, ui_collapsed=True)
labels = da.unique(segmented).compute()
labels
mask = da.from_array((segmented > 0).astype(np.uint8), chunks=beads.chunks)
mask
voxel_counts_s = dask_image.ndmeasure.sum(mask, segmented, labels).compute()
voxel_counts_s
voxel_counts = dask_image.ndmeasure.area(beads, segmented, labels).compute()
voxel_counts
voxel_counts - voxel_counts_s
def to_voxel_count_image(image_chunk):
    segmented = processing_pipeline(image_chunk)
    labels = np.unique(segmented)
    voxel_counts = dask_image.ndmeasure.area(image_chunk, segmented, labels)
    assert len(labels) == len(voxel_counts)
    result = np.zeros_like(image_chunk, np.uint16)
    for index in range(1, len(labels)):
        label = labels[index]
        result[segmented == label] = voxel_counts[index]
    return result                                  

voxel_count_image = da.map_blocks(to_voxel_count_image,
              beads,
              dtype=np.uint16)

voxel_count_image = voxel_count_image.compute()
voxel_count_image_rechunk = da.map_blocks(to_voxel_count_image,
              beads_rechunk,
              dtype=np.uint16)

voxel_count_image_rechunk = voxel_count_image_rechunk.compute()
voxel_count_max = da.maximum(voxel_count_image, voxel_count_image_rechunk)
chunk_labels = da.unique(segmented).compute()
max_chunk_labels = len(chunk_labels)
max_chunk_labels
def unique_labels(segmented,
                  segmented_rechunk,
                  voxel_count_image,
                  voxel_count_image_rechunk,
                  max_chunk_labels,
                  block_info=None):
    result = np.zeros_like(segmented)
    result[:,:,:] = block_info
    return result                                  

voxel_count_image = da.map_blocks(unique_labels,
              da.from_array(segmented, beads.chunk),
              da.from_array(segmented_rechunk, beads_rechunk.chunk),
              da.from_array(voxel_count_image_rechunk, beads_rechunk.chunk),
              max_chunk_labels
              dtype=np.uint16)

voxel_count_image = voxel_count_image.compute()
connected, num_features = dask_image.ndmeasure.label(voxel_count_max==23132)
connected = connected.astype(np.uint16)
view(connected)
view(beads, label_image=voxel_count_max, 
      label_image_blend=0.7,
      rotate=True,
      shadow=False,
      gradient_opacity=0.5)
gt = itk.imread('../data/segmented_bead_pack_GT.tif')
view(gt)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.09370677731673582, 'fraction_numerical': 0.01123789764868603, 'mean_word_length': 3.642857142857143, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '168722', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 1995, 'n_tokens_neox': 1931, 'n_words': 304, 'path': 'code/extra/01_b_bead_segmentation_distributed.ipynb', 'repo_name': 'dani-lbnl/SC20_pyHPC', 'size': '5784'}