1 2
import sys
2 2
import dask
3 2
import dask.array as da
4 2
from dask.diagnostics import ProgressBar
5 2
import warnings
6 2
import numpy as np
7 2
from numba import njit, prange
8 2
from edt import edt
9 2
import operator as op
10 2
from tqdm import tqdm
11 2
import scipy.ndimage as spim
12 2
import scipy.spatial as sptl
13 2
from collections import namedtuple
14 2
from scipy.signal import fftconvolve
15 2
from skimage.segmentation import clear_border, watershed
16 2
from skimage.morphology import ball, disk, square, cube, diamond, octahedron
17 2
from skimage.morphology import reconstruction
18 2
from porespy.tools import randomize_colors, fftmorphology
19 2
from porespy.tools import get_border, extend_slice, extract_subsection
20 2
from porespy.tools import _create_alias_map
21 2
from porespy.tools import ps_disk, ps_ball
22

23

24 2
def apply_padded(im, pad_width, func, pad_val=1, **kwargs):
25
    r"""
26
    Applies padding to an image before sending to ``func``, then extracts the
27
    result corresponding to the original image shape.
28

29
    Parameters
30
    ----------
31
    im : ND-image
32
        The image to which ``func`` should be applied
33
    pad_width : int or list of ints
34
        The amount of padding to apply to each axis.  Refer to ``numpy.pad``
35
        documentation for more details.
36
    pad_val : scalar
37
        The value to place into the padded voxels.  The default is 1 (or
38
        ``True``) which extends the pore space.
39
    func : function handle
40
        The function to apply to the padded image
41
    kwargs : additional keyword arguments
42
        All additional keyword arguments are collected and passed to ``func``.
43

44
    Notes
45
    -----
46
    A use case for this is when using ``skimage.morphology.skeletonize_3d``
47
    to ensure that the skeleton extends beyond the edges of the image.
48
    """
49 2
    padded = np.pad(im, pad_width=pad_width,
50
                    mode='constant', constant_values=pad_val)
51 2
    temp = func(padded, **kwargs)
52 2
    result = extract_subsection(im=temp, shape=im.shape)
53 2
    return result
54

55

56 2
def trim_small_clusters(im, size=1):
57
    r"""
58
    Remove isolated voxels or clusters of a given size or smaller
59

60
    Parameters
61
    ----------
62
    im : ND-array
63
        The binary image from which voxels are to be removed
64
    size : scalar
65
        The threshold size of clusters to trim.  As clusters with this many
66
        voxels or fewer will be trimmed.  The default is 1 so only single
67
        voxels are removed.
68

69
    Returns
70
    -------
71
    im : ND-image
72
        A copy of ``im`` with clusters of voxels smaller than the given
73
        ``size`` removed.
74

75
    """
76 0
    if im.ndim == 2:
77 0
        strel = disk(1)
78 0
    elif im.ndim == 3:
79 0
        strel = ball(1)
80
    else:
81 0
        raise Exception("Only 2D or 3D images are accepted")
82 0
    filtered_array = np.copy(im)
83 0
    labels, N = spim.label(filtered_array, structure=strel)
84 0
    id_sizes = np.array(spim.sum(im, labels, range(N + 1)))
85 0
    area_mask = id_sizes <= size
86 0
    filtered_array[area_mask[labels]] = 0
87 0
    return filtered_array
88

89

90 2
def hold_peaks(im, axis=-1):
91
    r"""
92
    Replaces each voxel with the highest value along the given axis
93

94
    Parameters
95
    ----------
96
    im : ND-image
97
        A greyscale image whose peaks are to
98
    """
99 0
    A = im
100 0
    B = np.swapaxes(A, axis, -1)
101 0
    updown = np.empty((*B.shape[:-1], B.shape[-1] + 1), B.dtype)
102 0
    updown[..., 0], updown[..., -1] = -1, -1
103 0
    np.subtract(B[..., 1:], B[..., :-1], out=updown[..., 1:-1])
104 0
    chnidx = np.where(updown)
105 0
    chng = updown[chnidx]
106 0
    (pkidx,) = np.where((chng[:-1] > 0) & (chng[1:] < 0) | (chnidx[-1][:-1] == 0))
107 0
    pkidx = (*map(op.itemgetter(pkidx), chnidx),)
108 0
    out = np.zeros_like(A)
109 0
    aux = out.swapaxes(axis, -1)
110 0
    aux[(*map(op.itemgetter(slice(1, None)), pkidx),)] = np.diff(B[pkidx])
111 0
    aux[..., 0] = B[..., 0]
112 0
    result = out.cumsum(axis=axis)
113 0
    return result
114

115

116 2
def distance_transform_lin(im, axis=0, mode="both"):
117
    r"""
118
    Replaces each void voxel with the linear distance to the nearest solid
119
    voxel along the specified axis.
120

121
    Parameters
122
    ----------
123
    im : ND-array
124
        The image of the porous material with ``True`` values indicating the
125
        void phase (or phase of interest)
126

127
    axis : int
128
        The direction along which the distance should be measured, the default
129
        is 0 (i.e. along the x-direction)
130

131
    mode : string
132
        Controls how the distance is measured.  Options are:
133

134
        'forward' - Distances are measured in the increasing direction along
135
        the specified axis
136

137
        'reverse' - Distances are measured in the reverse direction.
138
        *'backward'* is also accepted.
139

140
        'both' - Distances are calculated in both directions (by recursively
141
        calling itself), then reporting the minimum value of the two results.
142

143
    Returns
144
    -------
145
    image : ND-array
146
        A copy of ``im`` with each foreground voxel containing the distance to
147
        the nearest background along the specified axis.
148
    """
149 2
    if im.ndim != im.squeeze().ndim:
150 0
        warnings.warn(
151
            "Input image conains a singleton axis:"
152
            + str(im.shape)
153
            + " Reduce dimensionality with np.squeeze(im) to avoid"
154
            + " unexpected behavior."
155
        )
156 2
    if mode in ["backward", "reverse"]:
157 2
        im = np.flip(im, axis)
158 2
        im = distance_transform_lin(im=im, axis=axis, mode="forward")
159 2
        im = np.flip(im, axis)
160 2
        return im
161 2
    elif mode in ["both"]:
162 2
        im_f = distance_transform_lin(im=im, axis=axis, mode="forward")
163 2
        im_b = distance_transform_lin(im=im, axis=axis, mode="backward")
164 2
        return np.minimum(im_f, im_b)
165
    else:
166 2
        b = np.cumsum(im > 0, axis=axis)
167 2
        c = np.diff(b * (im == 0), axis=axis)
168 2
        d = np.minimum.accumulate(c, axis=axis)
169 2
        if im.ndim == 1:
170 0
            e = np.pad(d, pad_width=[1, 0], mode="constant", constant_values=0)
171 2
        elif im.ndim == 2:
172 2
            ax = [[[1, 0], [0, 0]], [[0, 0], [1, 0]]]
173 2
            e = np.pad(d, pad_width=ax[axis], mode="constant", constant_values=0)
174 2
        elif im.ndim == 3:
175 2
            ax = [
176
                [[1, 0], [0, 0], [0, 0]],
177
                [[0, 0], [1, 0], [0, 0]],
178
                [[0, 0], [0, 0], [1, 0]],
179
            ]
180 2
            e = np.pad(d, pad_width=ax[axis], mode="constant", constant_values=0)
181 2
        f = im * (b + e)
182 2
        return f
183

184

185 2
def snow_partitioning(im, dt=None, r_max=4, sigma=0.4, return_all=False,
186
                      mask=True, randomize=True):
187
    r"""
188
    Partitions the void space into pore regions using a marker-based watershed
189
    algorithm, with specially filtered peaks as markers.
190

191
    The SNOW network extraction algorithm (Sub-Network of an Over-segmented
192
    Watershed) was designed to handle to perculiarities of high porosity
193
    materials, but it applies well to other materials as well.
194

195
    Parameters
196
    ----------
197
    im : array_like
198
        A boolean image of the domain, with ``True`` indicating the pore space
199
        and ``False`` elsewhere.
200
    dt : array_like, optional
201
        The distance transform of the pore space.  This is done automatically
202
        if not provided, but if the distance transform has already been
203
        computed then supplying it can save some time.
204
    r_max : int
205
        The radius of the spherical structuring element to use in the Maximum
206
        filter stage that is used to find peaks.  The default is 4
207
    sigma : float
208
        The standard deviation of the Gaussian filter used in step 1.  The
209
        default is 0.4.  If 0 is given then the filter is not applied, which is
210
        useful if a distance transform is supplied as the ``im`` argument that
211
        has already been processed.
212
    return_all : boolean
213
        If set to ``True`` a named tuple is returned containing the original
214
        image, the distance transform, the filtered peaks, and the final
215
        pore regions.  The default is ``False``
216
    mask : boolean
217
        Apply a mask to the regions where the solid phase is.  Default is
218
        ``True``
219
    randomize : boolean
220
        If ``True`` (default), then the region colors will be randomized before
221
        returning.  This is helpful for visualizing otherwise neighboring
222
        regions have simlar coloring are are hard to distinguish.
223

224
    Returns
225
    -------
226
    image : ND-array
227
        An image the same shape as ``im`` with the void space partitioned into
228
        pores using a marker based watershed with the peaks found by the
229
        SNOW algorithm [1].
230

231
    Notes
232
    -----
233
    If ``return_all`` is ``True`` then a **named tuple** is returned containing
234
    all of the images used during the process.  They can be access as
235
    attriutes with the following names:
236

237
        * ``im``: The binary image of the void space
238
        * ``dt``: The distance transform of the image
239
        * ``peaks``: The peaks of the distance transform after applying the
240
        steps of the SNOW algorithm
241
        * ``regions``: The void space partitioned into pores using a marker
242
        based watershed with the peaks found by the SNOW algorithm
243

244
    References
245
    ----------
246
    [1] Gostick, J. "A versatile and efficient network extraction algorithm
247
    using marker-based watershed segmenation".  Physical Review E. (2017)
248

249
    """
250 2
    tup = namedtuple("results", field_names=["im", "dt", "peaks", "regions"])
251 2
    print("-" * 60)
252 2
    print("Beginning SNOW Algorithm")
253 2
    im_shape = np.array(im.shape)
254 2
    if im.dtype is not bool:
255 2
        print("Converting supplied image (im) to boolean")
256 2
        im = im > 0
257 2
    if dt is None:
258 2
        print("Peforming Distance Transform")
259 2
        if np.any(im_shape == 1):
260 2
            ax = np.where(im_shape == 1)[0][0]
261 2
            dt = edt(im.squeeze())
262 2
            dt = np.expand_dims(dt, ax)
263
        else:
264 2
            dt = edt(im)
265

266 2
    tup.im = im
267 2
    tup.dt = dt
268

269 2
    if sigma > 0:
270 2
        print("Applying Gaussian blur with sigma =", str(sigma))
271 2
        dt = spim.gaussian_filter(input=dt, sigma=sigma)
272

273 2
    peaks = find_peaks(dt=dt, r_max=r_max)
274 2
    print("Initial number of peaks: ", spim.label(peaks)[1])
275 2
    peaks = trim_saddle_points(peaks=peaks, dt=dt, max_iters=500)
276 2
    print("Peaks after trimming saddle points: ", spim.label(peaks)[1])
277 2
    peaks = trim_nearby_peaks(peaks=peaks, dt=dt)
278 2
    peaks, N = spim.label(peaks)
279 2
    print("Peaks after trimming nearby peaks: ", N)
280 2
    tup.peaks = peaks
281 2
    if mask:
282 2
        mask_solid = im > 0
283
    else:
284 0
        mask_solid = None
285 2
    regions = watershed(image=-dt, markers=peaks, mask=mask_solid)
286 2
    if randomize:
287 2
        regions = randomize_colors(regions)
288 2
    if return_all:
289 2
        tup.regions = regions
290 2
        return tup
291
    else:
292 2
        return regions
293

294

295 2
def snow_partitioning_n(im, r_max=4, sigma=0.4, return_all=True,
296
                        mask=True, randomize=False, alias=None):
297
    r"""
298
    This function partitions an imaging oontain an arbitrary number of phases
299
    into regions using a marker-based watershed segmentation. Its an extension
300
    of snow_partitioning function with all phases partitioned together.
301

302
    Parameters
303
    ----------
304
    im : ND-array
305
        Image of porous material where each phase is represented by unique
306
        integer starting from 1 (0's are ignored).
307
    r_max : scalar
308
        The radius of the spherical structuring element to use in the Maximum
309
        filter stage that is used to find peaks.  The default is 4.
310
    sigma : scalar
311
        The standard deviation of the Gaussian filter used.  The default is
312
        0.4. If 0 is given then the filter is not applied, which is useful if a
313
        distance transform is supplied as the ``im`` argument that has already
314
        been processed.
315
    return_all : boolean (default is False)
316
        If set to ``True`` a named tuple is returned containing the original
317
        image, the combined distance transform, list of each phase max label,
318
        and the final combined regions of all phases.
319
    mask : boolean (default is True)
320
        Apply a mask to the regions which are not under concern.
321
    randomize : boolean
322
        If ``True`` (default), then the region colors will be randomized before
323
        returning.  This is helpful for visualizing otherwise neighboring
324
        regions have similar coloring and are hard to distinguish.
325
    alias : dict (Optional)
326
        A dictionary that assigns unique image label to specific phases. For
327
        example {1: 'Solid'} will show all structural properties associated
328
        with label 1 as Solid phase properties. If ``None`` then default
329
        labelling will be used i.e {1: 'Phase1',..}.
330

331
    Returns
332
    -------
333
    An image the same shape as ``im`` with the all phases partitioned into
334
    regions using a marker based watershed with the peaks found by the
335
    SNOW algorithm [1].  If ``return_all`` is ``True`` then a **named tuple**
336
    is returned with the following attribute:
337

338
        * ``im`` : The actual image of the porous material
339
        * ``dt`` : The combined distance transform of the image
340
        * ``phase_max_label`` : The list of max label of each phase in order to
341
        distinguish between each other
342
        * ``regions`` : The partitioned regions of n phases using a marker
343
        based watershed with the peaks found by the SNOW algorithm
344

345
    References
346
    ----------
347
    [1] Gostick, J. "A versatile and efficient network extraction algorithm
348
    using marker-based watershed segmentation".  Physical Review E. (2017)
349

350
    [2] Khan, ZA et al. "Dual network extraction algorithm to investigate
351
    multiple transport processes in porous materials: Image-based modeling
352
    of pore and grain-scale processes".  Computers in Chemical Engineering.
353
    (2019)
354

355
    See Also
356
    ----------
357
    snow_partitioning
358

359
    Notes
360
    -----
361
    In principle it is possible to perform a distance transform on each
362
    phase separately, merge these into a single image, then apply the
363
    watershed only once. This, however, has been found to create edge artifacts
364
    between regions arising from the way watershed handles plateaus in the
365
    distance transform. To overcome this, this function applies the watershed
366
    to each of the distance transforms separately, then merges the segmented
367
    regions back into a single image.
368

369
    """
370
    # Get alias if provided by user
371 2
    al = _create_alias_map(im=im, alias=alias)
372
    # Perform snow on each phase and merge all segmentation and dt together
373 2
    phases_num = np.unique(im * 1)
374 2
    phases_num = np.trim_zeros(phases_num)
375 2
    combined_dt = 0
376 2
    combined_region = 0
377 2
    num = [0]
378 2
    for i, j in enumerate(phases_num):
379 2
        print("_" * 60)
380 2
        if alias is None:
381 2
            print("Processing Phase {}".format(j))
382
        else:
383 2
            print("Processing Phase {}".format(al[j]))
384 2
        phase_snow = snow_partitioning(
385
            im == j,
386
            dt=None,
387
            r_max=r_max,
388
            sigma=sigma,
389
            return_all=return_all,
390
            mask=mask,
391
            randomize=randomize,
392
        )
393 2
        combined_dt += phase_snow.dt
394 2
        phase_snow.regions *= phase_snow.im
395 2
        phase_snow.regions += num[i]
396 2
        phase_ws = phase_snow.regions * phase_snow.im
397 2
        phase_ws[phase_ws == num[i]] = 0
398 2
        combined_region += phase_ws
399 2
        num.append(np.amax(combined_region))
400 2
    if return_all:
401 2
        tup = namedtuple(
402
            "results", field_names=["im", "dt", "phase_max_label", "regions"]
403
        )
404 2
        tup.im = im
405 2
        tup.dt = combined_dt
406 2
        tup.phase_max_label = num[1:]
407 2
        tup.regions = combined_region
408 2
        return tup
409
    else:
410 0
        return combined_region
411

412

413 2
def find_peaks(dt, r_max=4, footprint=None, **kwargs):
414
    r"""
415
    Finds local maxima in the distance transform
416

417
    Parameters
418
    ----------
419
    dt : ND-array
420
        The distance transform of the pore space.  This may be calculated and
421
        filtered using any means desired.
422
    r_max : scalar
423
        The size of the structuring element used in the maximum filter.  This
424
        controls the localness of any maxima. The default is 4 voxels.
425
    footprint : ND-array
426
        Specifies the shape of the structuring element used to define the
427
        neighborhood when looking for peaks.  If ``None`` (the default) is
428
        specified then a spherical shape is used (or circular in 2D).
429

430
    Returns
431
    -------
432
    image : ND-array
433
        An array of booleans with ``True`` values at the location of any
434
        local maxima.
435

436
    Notes
437
    -----
438
    It is also possible ot the ``peak_local_max`` function from the
439
    ``skimage.feature`` module as follows:
440

441
    ``peaks = peak_local_max(image=dt, min_distance=r, exclude_border=0,
442
    indices=False)``
443

444
    The *skimage* function automatically uses a square structuring element
445
    which is significantly faster than using a circular or spherical element.
446
    """
447 2
    im = dt > 0
448 2
    if im.ndim != im.squeeze().ndim:
449 2
        warnings.warn("Input image conains a singleton axis:"
450
                      + str(im.shape)
451
                      + " Reduce dimensionality with np.squeeze(im) to avoid"
452
                      + " unexpected behavior.")
453 2
    if footprint is None:
454 2
        if im.ndim == 2:
455 2
            footprint = disk
456 2
        elif im.ndim == 3:
457 2
            footprint = ball
458
        else:
459 0
            raise Exception("only 2-d and 3-d images are supported")
460 2
    parallel = kwargs.pop('parallel', False)
461 2
    cores = kwargs.pop('cores', None)
462 2
    divs = kwargs.pop('cores', 2)
463 2
    if parallel:
464 2
        overlap = max(footprint(r_max).shape)
465 2
        peaks = chunked_func(func=find_peaks, overlap=overlap,
466
                             im_arg='dt', dt=dt, footprint=footprint,
467
                             cores=cores, divs=divs)
468
    else:
469 2
        mx = spim.maximum_filter(dt + 2 * (~im), footprint=footprint(r_max))
470 2
        peaks = (dt == mx) * im
471 2
    return peaks
472

473

474 2
def reduce_peaks(peaks):
475
    r"""
476
    Any peaks that are broad or elongated are replaced with a single voxel
477
    that is located at the center of mass of the original voxels.
478

479
    Parameters
480
    ----------
481
    peaks : ND-image
482
        An image containing ``True`` values indicating peaks in the distance
483
        transform
484

485
    Returns
486
    -------
487
    image : ND-array
488
        An array with the same number of isolated peaks as the original image,
489
        but fewer total ``True`` voxels.
490

491
    Notes
492
    -----
493
    The center of mass of a group of voxels is used as the new single voxel, so
494
    if the group has an odd shape (like a horse shoe), the new voxel may *not*
495
    lie on top of the original set.
496
    """
497 2
    if peaks.ndim == 2:
498 2
        strel = square
499
    else:
500 2
        strel = cube
501 2
    markers, N = spim.label(input=peaks, structure=strel(3))
502 2
    inds = spim.measurements.center_of_mass(
503
        input=peaks, labels=markers, index=np.arange(1, N + 1)
504
    )
505 2
    inds = np.floor(inds).astype(int)
506
    # Centroid may not be on old pixel, so create a new peaks image
507 2
    peaks_new = np.zeros_like(peaks, dtype=bool)
508 2
    peaks_new[tuple(inds.T)] = True
509 2
    return peaks_new
510

511

512 2
def trim_saddle_points(peaks, dt, max_iters=10, verbose=1):
513
    r"""
514
    Removes peaks that were mistakenly identified because they lied on a
515
    saddle or ridge in the distance transform that was not actually a true
516
    local peak.
517

518
    Parameters
519
    ----------
520
    peaks : ND-array
521
        A boolean image containing True values to mark peaks in the distance
522
        transform (``dt``)
523

524
    dt : ND-array
525
        The distance transform of the pore space for which the true peaks are
526
        sought.
527

528
    max_iters : int
529
        The maximum number of iterations to run while eroding the saddle
530
        points.  The default is 10, which is usually not reached; however,
531
        a warning is issued if the loop ends prior to removing all saddle
532
        points.
533

534
    Returns
535
    -------
536
    image : ND-array
537
        An image with fewer peaks than the input image
538

539
    References
540
    ----------
541
    [1] Gostick, J. "A versatile and efficient network extraction algorithm
542
    using marker-based watershed segmenation".  Physical Review E. (2017)
543

544
    """
545 2
    peaks = np.copy(peaks)
546 2
    if dt.ndim == 2:
547 2
        from skimage.morphology import square as cube
548
    else:
549 2
        from skimage.morphology import cube
550 2
    labels, N = spim.label(peaks)
551 2
    slices = spim.find_objects(labels)
552 2
    for i in range(N):
553 2
        s = extend_slice(s=slices[i], shape=peaks.shape, pad=10)
554 2
        peaks_i = labels[s] == i + 1
555 2
        dt_i = dt[s]
556 2
        im_i = dt_i > 0
557 2
        iters = 0
558 2
        peaks_dil = np.copy(peaks_i)
559 2
        while iters < max_iters:
560 2
            iters += 1
561 2
            peaks_dil = spim.binary_dilation(input=peaks_dil, structure=cube(3))
562 2
            peaks_max = peaks_dil * np.amax(dt_i * peaks_dil)
563 2
            peaks_extended = (peaks_max == dt_i) * im_i
564 2
            if np.all(peaks_extended == peaks_i):
565 2
                break  # Found a true peak
566 2
            elif np.sum(peaks_extended * peaks_i) == 0:
567 2
                peaks_i = False
568 2
                break  # Found a saddle point
569 2
        peaks[s] = peaks_i
570 2
        if iters >= max_iters and verbose:
571 0
            print(
572
                "Maximum number of iterations reached, consider "
573
                + "running again with a larger value of max_iters"
574
            )
575 2
    return peaks
576

577

578 2
def trim_nearby_peaks(peaks, dt):
579
    r"""
580
    Finds pairs of peaks that are nearer to each other than to the solid phase,
581
    and removes the peak that is closer to the solid.
582

583
    Parameters
584
    ----------
585
    peaks : ND-array
586
        A boolean image containing True values to mark peaks in the distance
587
        transform (``dt``)
588

589
    dt : ND-array
590
        The distance transform of the pore space for which the true peaks are
591
        sought.
592

593
    Returns
594
    -------
595
    image : ND-array
596
        An array the same size as ``peaks`` containing a subset of the peaks
597
        in the original image.
598

599
    Notes
600
    -----
601
    Each pair of peaks is considered simultaneously, so for a triplet of peaks
602
    each pair is considered.  This ensures that only the single peak that is
603
    furthest from the solid is kept.  No iteration is required.
604

605
    References
606
    ----------
607
    [1] Gostick, J. "A versatile and efficient network extraction algorithm
608
    using marker-based watershed segmenation".  Physical Review E. (2017)
609
    """
610 2
    peaks = np.copy(peaks)
611 2
    if dt.ndim == 2:
612 2
        from skimage.morphology import square as cube
613
    else:
614 2
        from skimage.morphology import cube
615 2
    peaks, N = spim.label(peaks, structure=cube(3))
616 2
    crds = spim.measurements.center_of_mass(peaks, labels=peaks,
617
                                            index=np.arange(1, N + 1))
618 2
    crds = np.vstack(crds).astype(int)  # Convert to numpy array of ints
619
    # Get distance between each peak as a distance map
620 2
    tree = sptl.cKDTree(data=crds)
621 2
    temp = tree.query(x=crds, k=2)
622 2
    nearest_neighbor = temp[1][:, 1]
623 2
    dist_to_neighbor = temp[0][:, 1]
624 2
    del temp, tree  # Free-up memory
625 2
    dist_to_solid = dt[tuple(crds.T)]  # Get distance to solid for each peak
626 2
    hits = np.where(dist_to_neighbor < dist_to_solid)[0]
627
    # Drop peak that is closer to the solid than it's neighbor
628 2
    drop_peaks = []
629 2
    for peak in hits:
630 2
        if dist_to_solid[peak] < dist_to_solid[nearest_neighbor[peak]]:
631 2
            drop_peaks.append(peak)
632
        else:
633 2
            drop_peaks.append(nearest_neighbor[peak])
634 2
    drop_peaks = np.unique(drop_peaks)
635
    # Remove peaks from image
636 2
    slices = spim.find_objects(input=peaks)
637 2
    for s in drop_peaks:
638 2
        peaks[slices[s]] = 0
639 2
    return peaks > 0
640

641

642 2
def find_disconnected_voxels(im, conn=None):
643
    r"""
644
    This identifies all pore (or solid) voxels that are not connected to the
645
    edge of the image.  This can be used to find blind pores, or remove
646
    artifacts such as solid phase voxels that are floating in space.
647

648
    Parameters
649
    ----------
650
    im : ND-image
651
        A Boolean image, with True values indicating the phase for which
652
        disconnected voxels are sought.
653
    conn : int
654
        For 2D the options are 4 and 8 for square and diagonal neighbors, while
655
        for the 3D the options are 6 and 26, similarily for square and diagonal
656
        neighbors.  The default is the maximum option.
657

658
    Returns
659
    -------
660
    image : ND-array
661
        An ND-image the same size as ``im``, with True values indicating
662
        voxels of the phase of interest (i.e. True values in the original
663
        image) that are not connected to the outer edges.
664

665
    Notes
666
    -----
667
    image : ND-array
668
        The returned array (e.g. ``holes``) be used to trim blind pores from
669
        ``im`` using: ``im[holes] = False``
670

671
    """
672 2
    if im.ndim != im.squeeze().ndim:
673 0
        warnings.warn(
674
            "Input image conains a singleton axis:"
675
            + str(im.shape)
676
            + " Reduce dimensionality with np.squeeze(im) to avoid"
677
            + " unexpected behavior."
678
        )
679 2
    if im.ndim == 2:
680 2
        if conn == 4:
681 2
            strel = disk(1)
682 2
        elif conn in [None, 8]:
683 2
            strel = square(3)
684
        else:
685 0
            raise Exception("Received conn is not valid")
686 2
    elif im.ndim == 3:
687 2
        if conn == 6:
688 2
            strel = ball(1)
689 2
        elif conn in [None, 26]:
690 2
            strel = cube(3)
691
        else:
692 0
            raise Exception("Received conn is not valid")
693 2
    labels, N = spim.label(input=im, structure=strel)
694 2
    holes = clear_border(labels=labels) > 0
695 2
    return holes
696

697

698 2
def fill_blind_pores(im, conn=None):
699
    r"""
700
    Fills all pores that are not connected to the edges of the image.
701

702
    Parameters
703
    ----------
704
    im : ND-array
705
        The image of the porous material
706

707
    Returns
708
    -------
709
    image : ND-array
710
        A version of ``im`` but with all the disconnected pores removed.
711
    conn : int
712
        For 2D the options are 4 and 8 for square and diagonal neighbors, while
713
        for the 3D the options are 6 and 26, similarily for square and diagonal
714
        neighbors.  The default is the maximum option.
715

716
    See Also
717
    --------
718
    find_disconnected_voxels
719

720
    """
721 2
    im = np.copy(im)
722 2
    holes = find_disconnected_voxels(im, conn=conn)
723 2
    im[holes] = False
724 2
    return im
725

726

727 2
def trim_floating_solid(im, conn=None):
728
    r"""
729
    Removes all solid that that is not attached to the edges of the image.
730

731
    Parameters
732
    ----------
733
    im : ND-array
734
        The image of the porous material
735
    conn : int
736
        For 2D the options are 4 and 8 for square and diagonal neighbors, while
737
        for the 3D the options are 6 and 26, similarily for square and diagonal
738
        neighbors.  The default is the maximum option.
739

740
    Returns
741
    -------
742
    image : ND-array
743
        A version of ``im`` but with all the disconnected solid removed.
744

745
    See Also
746
    --------
747
    find_disconnected_voxels
748

749
    """
750 2
    im = np.copy(im)
751 2
    holes = find_disconnected_voxels(~im, conn=conn)
752 2
    im[holes] = True
753 2
    return im
754

755

756 2
def trim_nonpercolating_paths(im, inlet_axis=0, outlet_axis=0,
757
                              inlets=None, outlets=None):
758
    r"""
759
    Removes all nonpercolating paths between specified edges
760

761
    This function is essential when performing transport simulations on an
762
    image, since image regions that do not span between the desired inlet and
763
    outlet do not contribute to the transport.
764

765
    Parameters
766
    ----------
767
    im : ND-array
768
        The image of the porous material with ```True`` values indicating the
769
        phase of interest
770
    inlet_axis : int
771
        Inlet axis of boundary condition. For three dimensional image the
772
        number ranges from 0 to 2. For two dimensional image the range is
773
        between 0 to 1. If ``inlets`` is given then this argument is ignored.
774
    outlet_axis : int
775
        Outlet axis of boundary condition. For three dimensional image the
776
        number ranges from 0 to 2. For two dimensional image the range is
777
        between 0 to 1. If ``outlets`` is given then this argument is ignored.
778
    inlets : ND-image (optional)
779
        A boolean mask indicating locations of inlets.  If this argument is
780
        supplied then ``inlet_axis`` is ignored.
781
    outlets : ND-image (optional)
782
        A boolean mask indicating locations of outlets. If this argument is
783
        supplied then ``outlet_axis`` is ignored.
784

785
    Returns
786
    -------
787
    image : ND-array
788
        A copy of ``im`` with all the nonpercolating paths removed
789

790
    See Also
791
    --------
792
    find_disconnected_voxels
793
    trim_floating_solid
794
    trim_blind_pores
795

796
    """
797 2
    if im.ndim != im.squeeze().ndim:
798 0
        warnings.warn(
799
            "Input image conains a singleton axis:"
800
            + str(im.shape)
801
            + " Reduce dimensionality with np.squeeze(im) to avoid"
802
            + " unexpected behavior."
803
        )
804 2
    im = trim_floating_solid(~im)
805 2
    labels = spim.label(~im)[0]
806 2
    if inlets is None:
807 2
        inlets = np.zeros_like(im, dtype=bool)
808 2
        if im.ndim == 3:
809 2
            if inlet_axis == 0:
810 2
                inlets[0, :, :] = True
811 2
            elif inlet_axis == 1:
812 2
                inlets[:, 0, :] = True
813 2
            elif inlet_axis == 2:
814 2
                inlets[:, :, 0] = True
815 2
        if im.ndim == 2:
816 2
            if inlet_axis == 0:
817 2
                inlets[0, :] = True
818 2
            elif inlet_axis == 1:
819 2
                inlets[:, 0] = True
820 2
    if outlets is None:
821 2
        outlets = np.zeros_like(im, dtype=bool)
822 2
        if im.ndim == 3:
823 2
            if outlet_axis == 0:
824 2
                outlets[-1, :, :] = True
825 2
            elif outlet_axis == 1:
826 2
                outlets[:, -1, :] = True
827 2
            elif outlet_axis == 2:
828 2
                outlets[:, :, -1] = True
829 2
        if im.ndim == 2:
830 2
            if outlet_axis == 0:
831 2
                outlets[-1, :] = True
832 2
            elif outlet_axis == 1:
833 2
                outlets[:, -1] = True
834 2
    IN = np.unique(labels * inlets)
835 2
    OUT = np.unique(labels * outlets)
836 2
    new_im = np.isin(labels, list(set(IN) ^ set(OUT)), invert=True)
837 2
    im[new_im == 0] = True
838 2
    return ~im
839

840

841 2
def trim_extrema(im, h, mode="maxima"):
842
    r"""
843
    Trims local extrema in greyscale values by a specified amount.
844

845
    This essentially decapitates peaks and/or floods valleys.
846

847
    Parameters
848
    ----------
849
    im : ND-array
850
        The image whose extrema are to be removed
851

852
    h : float
853
        The height to remove from each peak or fill in each valley
854

855
    mode : string {'maxima' | 'minima' | 'extrema'}
856
        Specifies whether to remove maxima or minima or both
857

858
    Returns
859
    -------
860
    image : ND-array
861
        A copy of the input image with all the peaks and/or valleys removed.
862

863
    Notes
864
    -----
865
    This function is referred to as **imhmax** or **imhmin** in Matlab.
866

867
    """
868 2
    result = im
869 2
    if mode in ["maxima", "extrema"]:
870 2
        result = reconstruction(seed=im - h, mask=im, method="dilation")
871 2
    elif mode in ["minima", "extrema"]:
872 2
        result = reconstruction(seed=im + h, mask=im, method="erosion")
873 2
    return result
874

875

876 2
def flood(im, regions=None, mode="max"):
877
    r"""
878
    Floods/fills each region in an image with a single value based on the
879
    specific values in that region.
880

881
    The ``mode`` argument is used to determine how the value is calculated.
882

883
    Parameters
884
    ----------
885
    im : array_like
886
        An image with isolated regions with numerical values in each voxel,
887
        and 0's elsewhere.
888
    regions : array_like
889
        An array the same shape as ``im`` with each region labeled.  If
890
        ``None`` is supplied (default) then ``scipy.ndimage.label`` is used
891
        with its default arguments.
892
    mode : string
893
        Specifies how to determine which value should be used to flood each
894
        region.  Options are:
895

896
        'maximum' - Floods each region with the local maximum in that region
897

898
        'minimum' - Floods each region the local minimum in that region
899

900
        'median' - Floods each region the local median in that region
901

902
        'mean' - Floods each region the local mean in that region
903

904
        'size' - Floods each region with the size of that region
905

906
    Returns
907
    -------
908
    image : ND-array
909
        A copy of ``im`` with new values placed in each forground voxel based
910
        on the ``mode``.
911

912
    See Also
913
    --------
914
    props_to_image
915

916
    """
917 2
    mask = im > 0
918 2
    if regions is None:
919 2
        labels, N = spim.label(mask)
920
    else:
921 0
        labels = np.copy(regions)
922 0
        N = labels.max()
923 2
    mode = "sum" if mode == "size" else mode
924 2
    mode = "maximum" if mode == "max" else mode
925 2
    mode = "minimum" if mode == "min" else mode
926 2
    if mode in ["mean", "median", "maximum", "minimum", "sum"]:
927 2
        f = getattr(spim, mode)
928 2
        vals = f(input=im, labels=labels, index=range(0, N + 1))
929 2
        im_flooded = vals[labels]
930 2
        im_flooded = im_flooded * mask
931
    else:
932 0
        raise Exception(mode + " is not a recognized mode")
933 2
    return im_flooded
934

935

936 2
def find_dt_artifacts(dt):
937
    r"""
938
    Finds points in a distance transform that are closer to wall than solid.
939

940
    These points could *potentially* be erroneously high since their distance
941
    values do not reflect the possibility that solid may have been present
942
    beyond the border of the image but lost by trimming.
943

944
    Parameters
945
    ----------
946
    dt : ND-array
947
        The distance transform of the phase of interest
948

949
    Returns
950
    -------
951
    image : ND-array
952
        An ND-array the same shape as ``dt`` with numerical values indicating
953
        the maximum amount of error in each volxel, which is found by
954
        subtracting the distance to nearest edge of image from the distance
955
        transform value. In other words, this is the error that would be found
956
        if there were a solid voxel lurking just beyond the nearest edge of
957
        the image.  Obviously, voxels with a value of zero have no error.
958

959
    """
960 2
    temp = np.ones(shape=dt.shape) * np.inf
961 2
    for ax in range(dt.ndim):
962 2
        dt_lin = distance_transform_lin(np.ones_like(temp, dtype=bool),
963
                                        axis=ax, mode="both")
964 2
        temp = np.minimum(temp, dt_lin)
965 2
    result = np.clip(dt - temp, a_min=0, a_max=np.inf)
966 2
    return result
967

968

969 2
def region_size(im):
970
    r"""
971
    Replace each voxel with size of region to which it belongs
972

973
    Parameters
974
    ----------
975
    im : ND-array
976
        Either a boolean image wtih ``True`` indicating the features of
977
        interest, in which case ``scipy.ndimage.label`` will be applied to
978
        find regions, or a greyscale image with integer values indicating
979
        regions.
980

981
    Returns
982
    -------
983
    image : ND-array
984
        A copy of ``im`` with each voxel value indicating the size of the
985
        region to which it belongs.  This is particularly useful for finding
986
        chord sizes on the image produced by ``apply_chords``.
987

988
    See Also
989
    --------
990
    flood
991
    """
992 2
    if im.dtype == bool:
993 2
        im = spim.label(im)[0]
994 2
    counts = np.bincount(im.flatten())
995 2
    counts[0] = 0
996 2
    chords = counts[im]
997 2
    return chords
998

999

1000 2
def apply_chords(im, spacing=1, axis=0, trim_edges=True, label=False):
1001
    r"""
1002
    Adds chords to the void space in the specified direction.  The chords are
1003
    separated by 1 voxel plus the provided spacing.
1004

1005
    Parameters
1006
    ----------
1007
    im : ND-array
1008
        An image of the porous material with void marked as ``True``.
1009

1010
    spacing : int
1011
        Separation between chords.  The default is 1 voxel.  This can be
1012
        decreased to 0, meaning that the chords all touch each other, which
1013
        automatically sets to the ``label`` argument to ``True``.
1014

1015
    axis : int (default = 0)
1016
        The axis along which the chords are drawn.
1017

1018
    trim_edges : bool (default = ``True``)
1019
        Whether or not to remove chords that touch the edges of the image.
1020
        These chords are artifically shortened, so skew the chord length
1021
        distribution.
1022

1023
    label : bool (default is ``False``)
1024
        If ``True`` the chords in the returned image are each given a unique
1025
        label, such that all voxels lying on the same chord have the same
1026
        value.  This is automatically set to ``True`` if spacing is 0, but is
1027
        ``False`` otherwise.
1028

1029
    Returns
1030
    -------
1031
    image : ND-array
1032
        A copy of ``im`` with non-zero values indicating the chords.
1033

1034
    See Also
1035
    --------
1036
    apply_chords_3D
1037

1038
    """
1039 2
    if im.ndim != im.squeeze().ndim:
1040 0
        warnings.warn(
1041
            "Input image conains a singleton axis:"
1042
            + str(im.shape)
1043
            + " Reduce dimensionality with np.squeeze(im) to avoid"
1044
            + " unexpected behavior."
1045
        )
1046 2
    if spacing < 0:
1047 2
        raise Exception("Spacing cannot be less than 0")
1048 2
    if spacing == 0:
1049 0
        label = True
1050 2
    result = np.zeros(im.shape, dtype=int)  # Will receive chords at end
1051 2
    slxyz = [slice(None, None, spacing * (axis != i) + 1) for i in [0, 1, 2]]
1052 2
    slices = tuple(slxyz[: im.ndim])
1053 2
    s = [[0, 1, 0], [0, 1, 0], [0, 1, 0]]  # Straight-line structuring element
1054 2
    if im.ndim == 3:  # Make structuring element 3D if necessary
1055 2
        s = np.pad(np.atleast_3d(s), pad_width=((0, 0), (0, 0), (1, 1)),
1056
                   mode="constant", constant_values=0)
1057 2
    im = im[slices]
1058 2
    s = np.swapaxes(s, 0, axis)
1059 2
    chords = spim.label(im, structure=s)[0]
1060 2
    if trim_edges:  # Label on border chords will be set to 0
1061 2
        chords = clear_border(chords)
1062 2
    result[slices] = chords  # Place chords into empty image created at top
1063 2
    if label is False:  # Remove label if not requested
1064 2
        result = result > 0
1065 2
    return result
1066

1067

1068 2
def apply_chords_3D(im, spacing=0, trim_edges=True):
1069
    r"""
1070
    Adds chords to the void space in all three principle directions.  The
1071
    chords are seprated by 1 voxel plus the provided spacing.  Chords in the X,
1072
    Y and Z directions are labelled 1, 2 and 3 resepctively.
1073

1074
    Parameters
1075
    ----------
1076
    im : ND-array
1077
        A 3D image of the porous material with void space marked as True.
1078

1079
    spacing : int (default = 0)
1080
        Chords are automatically separed by 1 voxel on all sides, and this
1081
        argument increases the separation.
1082

1083
    trim_edges : bool (default is ``True``)
1084
        Whether or not to remove chords that touch the edges of the image.
1085
        These chords are artifically shortened, so skew the chord length
1086
        distribution
1087

1088
    Returns
1089
    -------
1090
    image : ND-array
1091
        A copy of ``im`` with values of 1 indicating x-direction chords,
1092
        2 indicating y-direction chords, and 3 indicating z-direction chords.
1093

1094
    Notes
1095
    -----
1096
    The chords are separated by a spacing of at least 1 voxel so that tools
1097
    that search for connected components, such as ``scipy.ndimage.label`` can
1098
    detect individual chords.
1099

1100
    See Also
1101
    --------
1102
    apply_chords
1103

1104
    """
1105 2
    if im.ndim != im.squeeze().ndim:
1106 0
        warnings.warn(
1107
            "Input image conains a singleton axis:"
1108
            + str(im.shape)
1109
            + " Reduce dimensionality with np.squeeze(im) to avoid"
1110
            + " unexpected behavior."
1111
        )
1112 2
    if im.ndim < 3:
1113 0
        raise Exception("Must be a 3D image to use this function")
1114 2
    if spacing < 0:
1115 0
        raise Exception("Spacing cannot be less than 0")
1116 2
    ch = np.zeros_like(im, dtype=int)
1117 2
    ch[:, :: 4 + 2 * spacing, :: 4 + 2 * spacing] = 1       # X-direction
1118 2
    ch[:: 4 + 2 * spacing, :, 2::4 + 2 * spacing] = 2     # Y-direction
1119 2
    ch[2::4 + 2 * spacing, 2::4 + 2 * spacing, :] = 3   # Z-direction
1120 2
    chords = ch * im
1121 2
    if trim_edges:
1122 2
        temp = clear_border(spim.label(chords > 0)[0]) > 0
1123 2
        chords = temp * chords
1124 2
    return chords
1125

1126

1127 2
def local_thickness(im, sizes=25, mode="hybrid", **kwargs):
1128
    r"""
1129
    For each voxel, this function calculates the radius of the largest sphere
1130
    that both engulfs the voxel and fits entirely within the foreground.
1131

1132
    This is not the same as a simple distance transform, which finds the
1133
    largest sphere that could be *centered* on each voxel.
1134

1135
    Parameters
1136
    ----------
1137
    im : array_like
1138
        A binary image with the phase of interest set to True
1139

1140
    sizes : array_like or scalar
1141
        The sizes to invade.  If a list of values of provided they are used
1142
        directly.  If a scalar is provided then that number of points spanning
1143
        the min and max of the distance transform are used.
1144

1145
    mode : string
1146
        Controls with method is used to compute the result.  Options are:
1147

1148
        'hybrid' - (default) Performs a distance tranform of the void space,
1149
        thresholds to find voxels larger than ``sizes[i]``, trims the resulting
1150
        mask if ``access_limitations`` is ``True``, then dilates it using the
1151
        efficient fft-method to obtain the non-wetting fluid configuration.
1152

1153
        'dt' - Same as 'hybrid', except uses a second distance transform,
1154
        relative to the thresholded mask, to find the invading fluid
1155
        configuration.  The choice of 'dt' or 'hybrid' depends on speed, which
1156
        is system and installation specific.
1157

1158
        'mio' - Using a single morphological image opening step to obtain the
1159
        invading fluid confirguration directly, *then* trims if
1160
        ``access_limitations`` is ``True``.  This method is not ideal and is
1161
        included mostly for comparison purposes.
1162

1163
    Returns
1164
    -------
1165
    image : ND-array
1166
        A copy of ``im`` with the pore size values in each voxel
1167

1168
    See Also
1169
    --------
1170
    porosimetry
1171

1172
    Notes
1173
    -----
1174
    The term *foreground* is used since this function can be applied to both
1175
    pore space or the solid, whichever is set to ``True``.
1176

1177
    This function is identical to ``porosimetry`` with ``access_limited`` set
1178
    to ``False``.
1179

1180
    The way local thickness is found in PoreSpy differs from the traditional
1181
    method (i.e. `used in ImageJ <https://imagej.net/Local_Thickness>`_).
1182
    Our approach is probably slower, but it allows for the same code to be
1183
    used for ``local_thickness`` and ``porosimetry``, since we can 'trim'
1184
    invaded regions that are not connected to the inlets in the ``porosimetry``
1185
    function.  This is not needed in ``local_thickness`` however.
1186

1187
    """
1188 2
    im_new = porosimetry(im=im, sizes=sizes, access_limited=False, mode=mode,
1189
                         **kwargs)
1190 2
    return im_new
1191

1192

1193 2
def porosimetry(im, sizes=25, inlets=None, access_limited=True, mode='hybrid',
1194
                fft=True, **kwargs):
1195
    r"""
1196
    Performs a porosimetry simulution on an image
1197

1198
    Parameters
1199
    ----------
1200
    im : ND-array
1201
        An ND image of the porous material containing ``True`` values in the
1202
        pore space.
1203

1204
    sizes : array_like or scalar
1205
        The sizes to invade.  If a list of values of provided they are used
1206
        directly.  If a scalar is provided then that number of points spanning
1207
        the min and max of the distance transform are used.
1208

1209
    inlets : ND-array, boolean
1210
        A boolean mask with ``True`` values indicating where the invasion
1211
        enters the image.  By default all faces are considered inlets,
1212
        akin to a mercury porosimetry experiment.  Users can also apply
1213
        solid boundaries to their image externally before passing it in,
1214
        allowing for complex inlets like circular openings, etc.  This argument
1215
        is only used if ``access_limited`` is ``True``.
1216

1217
    access_limited : Boolean
1218
        This flag indicates if the intrusion should only occur from the
1219
        surfaces (``access_limited`` is ``True``, which is the default), or
1220
        if the invading phase should be allowed to appear in the core of
1221
        the image.  The former simulates experimental tools like mercury
1222
        intrusion porosimetry, while the latter is useful for comparison
1223
        to gauge the extent of shielding effects in the sample.
1224

1225
    mode : string
1226
        Controls with method is used to compute the result.  Options are:
1227

1228
        'hybrid' - (default) Performs a distance tranform of the void space,
1229
        thresholds to find voxels larger than ``sizes[i]``, trims the resulting
1230
        mask if ``access_limitations`` is ``True``, then dilates it using the
1231
        efficient fft-method to obtain the non-wetting fluid configuration.
1232

1233
        'dt' - Same as 'hybrid', except uses a second distance transform,
1234
        relative to the thresholded mask, to find the invading fluid
1235
        configuration.  The choice of 'dt' or 'hybrid' depends on speed, which
1236
        is system and installation specific.
1237

1238
        'mio' - Using a single morphological image opening step to obtain the
1239
        invading fluid confirguration directly, *then* trims if
1240
        ``access_limitations`` is ``True``.  This method is not ideal and is
1241
        included mostly for comparison purposes.  The morphological operations
1242
        are done using fft-based method implementations.
1243

1244
    fft : boolean (default is ``True``)
1245
        Indicates whether to use the ``fftmorphology`` function in
1246
        ``porespy.filters`` or to use the standard morphology functions in
1247
        ``scipy.ndimage``.  Always use ``fft=True`` unless you have a good
1248
        reason not to.
1249

1250
    Returns
1251
    -------
1252
    image : ND-array
1253
        A copy of ``im`` with voxel values indicating the sphere radius at
1254
        which it becomes accessible from the inlets.  This image can be used
1255
        to find invading fluid configurations as a function of applied
1256
        capillary pressure by applying a boolean comparison:
1257
        ``inv_phase = im > r`` where ``r`` is the radius (in voxels) of the
1258
        invading sphere.  Of course, ``r`` can be converted to capillary
1259
        pressure using a preferred model.
1260

1261
    Notes
1262
    -----
1263
    There are many ways to perform this filter, and PoreSpy offers 3, which
1264
    users can choose between via the ``mode`` argument.  These methods all
1265
    work in a similar way by finding which foreground voxels can accomodate
1266
    a sphere of a given radius, then repeating for smaller radii.
1267

1268
    See Also
1269
    --------
1270
    local_thickness
1271

1272
    """
1273 2
    if im.ndim != im.squeeze().ndim:
1274 0
        warnings.warn("Input image contains a singleton axis:"
1275
                      + str(im.shape)
1276
                      + " Reduce dimensionality with np.squeeze(im) to avoid"
1277
                      + " unexpected behavior.")
1278

1279 2
    dt = edt(im > 0)
1280

1281 2
    if inlets is None:
1282 2
        inlets = get_border(im.shape, mode="faces")
1283

1284 2
    if isinstance(sizes, int):
1285 2
        sizes = np.logspace(start=np.log10(np.amax(dt)), stop=0, num=sizes)
1286
    else:
1287 2
        sizes = np.unique(sizes)[-1::-1]
1288

1289 2
    if im.ndim == 2:
1290 2
        strel = ps_disk
1291
    else:
1292 2
        strel = ps_ball
1293

1294
    # Parse kwargs for any parallelization arguments
1295 2
    parallel = kwargs.pop('parallel', False)
1296 2
    cores = kwargs.pop('cores', None)
1297 2
    divs = kwargs.pop('divs', 2)
1298

1299 2
    if mode == "mio":
1300 2
        pw = int(np.floor(dt.max()))
1301 2
        impad = np.pad(im, mode="symmetric", pad_width=pw)
1302 2
        inlets = np.pad(inlets, mode="symmetric", pad_width=pw)
1303
        # sizes = np.unique(np.around(sizes, decimals=0).astype(int))[-1::-1]
1304 2
        imresults = np.zeros(np.shape(impad))
1305 2
        with tqdm(sizes) as pbar:
1306 2
            for r in sizes:
1307 2
                pbar.update()
1308 2
                if parallel:
1309 2
                    imtemp = chunked_func(func=spim.binary_erosion,
1310
                                          input=impad, structure=strel(r),
1311
                                          overlap=int(2*r) + 1,
1312
                                          cores=cores, divs=divs)
1313 2
                elif fft:
1314 2
                    imtemp = fftmorphology(impad, strel(r), mode="erosion")
1315
                else:
1316 2
                    imtemp = spim.binary_erosion(input=impad,
1317
                                                 structure=strel(r))
1318 2
                if access_limited:
1319 2
                    imtemp = trim_disconnected_blobs(imtemp, inlets)
1320 2
                if parallel:
1321 2
                    imtemp = chunked_func(func=spim.binary_dilation,
1322
                                          input=imtemp, structure=strel(r),
1323
                                          overlap=int(2*r) + 1,
1324
                                          cores=cores, divs=divs)
1325 2
                elif fft:
1326 2
                    imtemp = fftmorphology(imtemp, strel(r), mode="dilation")
1327
                else:
1328 2
                    imtemp = spim.binary_dilation(input=imtemp,
1329
                                                  structure=strel(r))
1330 2
                if np.any(imtemp):
1331 2
                    imresults[(imresults == 0) * imtemp] = r
1332 2
        imresults = extract_subsection(imresults, shape=im.shape)
1333 2
    elif mode == "dt":
1334 2
        imresults = np.zeros(np.shape(im))
1335 2
        with tqdm(sizes) as pbar:
1336 2
            for r in sizes:
1337 2
                pbar.update()
1338 2
                imtemp = dt >= r
1339 2
                if access_limited:
1340 2
                    imtemp = trim_disconnected_blobs(imtemp, inlets)
1341 2
                if np.any(imtemp):
1342 2
                    imtemp = edt(~imtemp) < r
1343 2
                    imresults[(imresults == 0) * imtemp] = r
1344 2
    elif mode == "hybrid":
1345 2
        imresults = np.zeros(np.shape(im))
1346 2
        with tqdm(sizes) as pbar:
1347 2
            for r in sizes:
1348 2
                pbar.update()
1349 2
                imtemp = dt >= r
1350 2
                if access_limited:
1351 2
                    imtemp = trim_disconnected_blobs(imtemp, inlets)
1352 2
                if np.any(imtemp):
1353 2
                    if parallel:
1354 0
                        imtemp = chunked_func(func=spim.binary_dilation,
1355
                                              input=imtemp, structure=strel(r),
1356
                                              overlap=int(2*r) + 1,
1357
                                              cores=cores, divs=divs)
1358 2
                    elif fft:
1359 2
                        imtemp = fftmorphology(imtemp, strel(r),
1360
                                               mode="dilation")
1361
                    else:
1362 2
                        imtemp = spim.binary_dilation(input=imtemp,
1363
                                                      structure=strel(r))
1364 2
                    imresults[(imresults == 0) * imtemp] = r
1365
    else:
1366 0
        raise Exception("Unrecognized mode " + mode)
1367 2
    return imresults
1368

1369

1370 2
def trim_disconnected_blobs(im, inlets, strel=None):
1371
    r"""
1372
    Removes foreground voxels not connected to specified inlets
1373

1374
    Parameters
1375
    ----------
1376
    im : ND-array
1377
        The image containing the blobs to be trimmed
1378
    inlets : ND-array or tuple of indices
1379
        The locations of the inlets.  Can either be a boolean mask the same
1380
        shape as ``im``, or a tuple of indices such as that returned by the
1381
        ``where`` function.  Any voxels *not* connected directly to
1382
        the inlets will be trimmed.
1383
    strel : ND-array
1384
        The neighborhood over which connectivity should be checked. It must
1385
        be symmetric and the same dimensionality as the image.  It is passed
1386
        directly to the ``scipy.ndimage.label`` function as the ``structure``
1387
        argument so refer to that docstring for additional info.
1388

1389
    Returns
1390
    -------
1391
    image : ND-array
1392
        An array of the same shape as ``im``, but with all foreground
1393
        voxels not connected to the ``inlets`` removed.
1394
    """
1395 2
    if type(inlets) == tuple:
1396 0
        temp = np.copy(inlets)
1397 0
        inlets = np.zeros_like(im, dtype=bool)
1398 0
        inlets[temp] = True
1399 2
    elif (inlets.shape == im.shape) and (inlets.max() == 1):
1400 2
        inlets = inlets.astype(bool)
1401
    else:
1402 0
        raise Exception("inlets not valid, refer to docstring for info")
1403 2
    if im.ndim == 3:
1404 2
        strel = cube
1405
    else:
1406 2
        strel = square
1407 2
    labels = spim.label(inlets + (im > 0), structure=strel(3))[0]
1408 2
    keep = np.unique(labels[inlets])
1409 2
    keep = keep[keep > 0]
1410 2
    if len(keep) > 0:
1411 2
        im2 = np.reshape(np.in1d(labels, keep), newshape=im.shape)
1412
    else:
1413 0
        im2 = np.zeros_like(im)
1414 2
    im2 = im2 * im
1415 2
    return im2
1416

1417

1418 2
def _get_axial_shifts(ndim=2, include_diagonals=False):
1419
    r"""
1420
    Helper function to generate the axial shifts that will be performed on
1421
    the image to identify bordering pixels/voxels
1422
    """
1423 2
    if ndim == 2:
1424 2
        if include_diagonals:
1425 2
            neighbors = square(3)
1426
        else:
1427 2
            neighbors = diamond(1)
1428 2
        neighbors[1, 1] = 0
1429 2
        x, y = np.where(neighbors)
1430 2
        x -= 1
1431 2
        y -= 1
1432 2
        return np.vstack((x, y)).T
1433
    else:
1434 2
        if include_diagonals:
1435 2
            neighbors = cube(3)
1436
        else:
1437 2
            neighbors = octahedron(1)
1438 2
        neighbors[1, 1, 1] = 0
1439 2
        x, y, z = np.where(neighbors)
1440 2
        x -= 1
1441 2
        y -= 1
1442 2
        z -= 1
1443 2
        return np.vstack((x, y, z)).T
1444

1445

1446 2
def _make_stack(im, include_diagonals=False):
1447
    r"""
1448
    Creates a stack of images with one extra dimension to the input image
1449
    with length equal to the number of borders to search + 1.
1450
    Image is rolled along the axial shifts so that the border pixel is
1451
    overlapping the original pixel. First image in stack is the original.
1452
    Stacking makes direct vectorized array comparisons possible.
1453
    """
1454 2
    ndim = len(np.shape(im))
1455 2
    axial_shift = _get_axial_shifts(ndim, include_diagonals)
1456 2
    if ndim == 2:
1457 2
        stack = np.zeros([np.shape(im)[0], np.shape(im)[1], len(axial_shift) + 1])
1458 2
        stack[:, :, 0] = im
1459 2
        for i in range(len(axial_shift)):
1460 2
            ax0, ax1 = axial_shift[i]
1461 2
            temp = np.roll(np.roll(im, ax0, 0), ax1, 1)
1462 2
            stack[:, :, i + 1] = temp
1463 2
        return stack
1464 2
    elif ndim == 3:
1465 2
        stack = np.zeros(
1466
            [np.shape(im)[0], np.shape(im)[1], np.shape(im)[2], len(axial_shift) + 1]
1467
        )
1468 2
        stack[:, :, :, 0] = im
1469 2
        for i in range(len(axial_shift)):
1470 2
            ax0, ax1, ax2 = axial_shift[i]
1471 2
            temp = np.roll(np.roll(np.roll(im, ax0, 0), ax1, 1), ax2, 2)
1472 2
            stack[:, :, :, i + 1] = temp
1473 2
        return stack
1474

1475

1476 2
def nphase_border(im, include_diagonals=False):
1477
    r"""
1478
    Identifies the voxels in regions that border *N* other regions.
1479

1480
    Useful for finding triple-phase boundaries.
1481

1482
    Parameters
1483
    ----------
1484
    im : ND-array
1485
        An ND image of the porous material containing discrete values in the
1486
        pore space identifying different regions. e.g. the result of a
1487
        snow-partition
1488

1489
    include_diagonals : boolean
1490
        When identifying bordering pixels (2D) and voxels (3D) include those
1491
        shifted along more than one axis
1492

1493
    Returns
1494
    -------
1495
    image : ND-array
1496
        A copy of ``im`` with voxel values equal to the number of uniquely
1497
        different bordering values
1498
    """
1499 2
    if im.ndim != im.squeeze().ndim:
1500 0
        warnings.warn(
1501
            "Input image conains a singleton axis:"
1502
            + str(im.shape)
1503
            + " Reduce dimensionality with np.squeeze(im) to avoid"
1504
            + " unexpected behavior."
1505
        )
1506
    # Get dimension of image
1507 2
    ndim = len(np.shape(im))
1508 2
    if ndim not in [2, 3]:
1509
        raise NotImplementedError("Function only works for 2d and 3d images")
1510
    # Pad image to handle edges
1511 2
    im = np.pad(im, pad_width=1, mode="edge")
1512
    # Stack rolled images for each neighbor to be inspected
1513 2
    stack = _make_stack(im, include_diagonals)
1514
    # Sort the stack along the last axis
1515 2
    stack.sort()
1516 2
    out = np.ones_like(im)
1517
    # Run through stack recording when neighbor id changes
1518
    # Number of changes is number of unique bordering regions
1519 2
    for k in range(np.shape(stack)[ndim])[1:]:
1520 2
        if ndim == 2:
1521 2
            mask = stack[:, :, k] != stack[:, :, k - 1]
1522 2
        elif ndim == 3:
1523 2
            mask = stack[:, :, :, k] != stack[:, :, :, k - 1]
1524 2
        out += mask
1525
    # Un-pad
1526 2
    if ndim == 2:
1527 2
        return out[1:-1, 1:-1].copy()
1528
    else:
1529 2
        return out[1:-1, 1:-1, 1:-1].copy()
1530

1531

1532 2
def prune_branches(skel, branch_points=None, iterations=1, **kwargs):
1533
    r"""
1534
    Removes all dangling ends or tails of a skeleton.
1535

1536
    Parameters
1537
    ----------
1538
    skel : ND-image
1539
        A image of a full or partial skeleton from which the tails should be
1540
        trimmed.
1541

1542
    branch_points : ND-image, optional
1543
        An image the same size ``skel`` with True values indicating the branch
1544
        points of the skeleton.  If this is not provided it is calculated
1545
        automatically.
1546

1547
    Returns
1548
    -------
1549
    An ND-image containing the skeleton with tails removed.
1550

1551
    """
1552 2
    skel = skel > 0
1553 2
    if skel.ndim == 2:
1554 0
        from skimage.morphology import square as cube
1555
    else:
1556 2
        from skimage.morphology import cube
1557 2
    parallel = kwargs.pop('parallel', False)
1558 2
    divs = kwargs.pop('divs', 2)
1559 2
    cores = kwargs.pop('cores', None)
1560
    # Create empty image to house results
1561 2
    im_result = np.zeros_like(skel)
1562
    # If branch points are not supplied, attempt to find them
1563 2
    if branch_points is None:
1564 2
        branch_points = spim.convolve(skel * 1.0, weights=cube(3)) > 3
1565 2
        branch_points = branch_points * skel
1566
    # Store original branch points before dilating
1567 2
    pts_orig = branch_points
1568
    # Find arcs of skeleton by deleting branch points
1569 2
    arcs = skel * (~branch_points)
1570
    # Label arcs
1571 2
    arc_labels = spim.label(arcs, structure=cube(3))[0]
1572
    # Dilate branch points so they overlap with the arcs
1573 2
    if parallel:
1574 2
        branch_points = chunked_func(func=spim.binary_dilation,
1575
                                     input=branch_points, structure=cube(3),
1576
                                     overlap=3, divs=divs, cores=cores)
1577
    else:
1578 2
        branch_points = spim.binary_dilation(branch_points, structure=cube(3))
1579 2
    pts_labels = spim.label(branch_points, structure=cube(3))[0]
1580
    # Now scan through each arc to see if it's connected to two branch points
1581 2
    slices = spim.find_objects(arc_labels)
1582 2
    label_num = 0
1583 2
    for s in slices:
1584 0
        label_num += 1
1585
        # Find branch point labels the overlap current arc
1586 0
        hits = pts_labels[s] * (arc_labels[s] == label_num)
1587
        # If image contains 2 branch points, then it's not a tail.
1588 0
        if len(np.unique(hits)) == 3:
1589 0
            im_result[s] += arc_labels[s] == label_num
1590
    # Add missing branch points back to arc image to make complete skeleton
1591 2
    im_result += skel * pts_orig
1592 2
    if iterations > 1:
1593 2
        iterations -= 1
1594 2
        im_temp = np.copy(im_result)
1595 2
        im_result = prune_branches(skel=im_result,
1596
                                   branch_points=None,
1597
                                   iterations=iterations,
1598
                                   parallel=parallel,
1599
                                   divs=divs, cores=cores)
1600 2
        if np.all(im_temp == im_result):
1601 2
            iterations = 0
1602 2
    return im_result
1603

1604

1605 2
def chunked_func(func,
1606
                 overlap=None,
1607
                 divs=2,
1608
                 cores=None,
1609
                 im_arg=["input", "image", "im"],
1610
                 strel_arg=["strel", "structure", "footprint"],
1611
                 **kwargs):
1612
    r"""
1613
    Performs the specfied operation "chunk-wise" in parallel using dask
1614

1615
    This can be used to save memory by doing one chunk at a time (``cores=1``)
1616
    or to increase computation speed by spreading the work across multiple
1617
    cores (e.g. ``cores=8``)
1618

1619
    This function can be used with any operation that applies a structuring
1620
    element of some sort, since this implies that the operation is local
1621
    and can be chunked.
1622

1623
    Parameters
1624
    ----------
1625
    func : function handle
1626
        The function which should be applied to each chunk, such as
1627
        ``spipy.ndimage.binary_dilation``.
1628
    overlap : scalar or list of scalars, optional
1629
        The amount of overlap to include when dividing up the image.  This
1630
        value will almost always be the size (i.e. diameter) of the
1631
        structuring element. If not specified then the amount of overlap is
1632
        inferred from the size of the structuring element, in which case the
1633
        ``strel_arg`` must be specified.
1634
    divs : scalar or list of scalars (default = [2, 2, 2])
1635
        The number of chunks to divide the image into in each direction.  The
1636
        default is 2 chunks in each direction, resulting in a quartering of
1637
        the image and 8 total chunks (in 3D).  A scalar is interpreted as
1638
        applying to all directions, while a list of scalars is interpreted
1639
        as applying to each individual direction.
1640
    cores : scalar
1641
        The number of cores which should be used.  By default, all cores will
1642
        be used, or as many are needed for the given number of chunks, which
1643
        ever is smaller.
1644
    im_arg : string
1645
        The keyword used by ``func`` for the image to be operated on.  By
1646
        default this function will look for ``image``, ``input``, and ``im``
1647
        which are commonly used by *scipy.ndimage* and *skimage*.
1648
    strel_arg : string
1649
        The keyword used by ``func`` for the structuring element to apply.
1650
        This is only needed if ``overlap`` is not specified. By default this
1651
        function will look for ``strel``, ``structure``, and ``footprint``
1652
        which are commonly used by *scipy.ndimage* and *skimage*.
1653
    kwargs : additional keyword arguments
1654
        All other arguments are passed to ``func`` as keyword arguments. Note
1655
        that PoreSpy will fetch the image from this list of keywords using the
1656
        value provided to ``im_arg``.
1657

1658
    Returns
1659
    -------
1660
    result : ND-image
1661
        An image the same size as the input image, with the specified filter
1662
        applied as though done on a single large image.  There should be *no*
1663
        difference.
1664

1665
    Notes
1666
    -----
1667
    This function divides the image into the specified number of chunks, but
1668
    also applies a padding to each chunk to create an overlap with neighboring
1669
    chunks.  This way the operation does not have any edge artifacts. The
1670
    amount of padding is usually equal to the radius of the structuring
1671
    element but some functions do not use one, such as the distance transform
1672
    and Gaussian blur.  In these cases the user can specify ``overlap``.
1673

1674
    See Also
1675
    --------
1676
    skikit-image.util.apply_parallel
1677

1678
    Examples
1679
    --------
1680
    >>> import scipy.ndimage as spim
1681
    >>> import porespy as ps
1682
    >>> from skimage.morphology import ball
1683
    >>> im = ps.generators.blobs(shape=[100, 100, 100])
1684
    >>> f = spim.binary_dilation
1685
    >>> im2 = ps.filters.chunked_func(func=f, overlap=7, im_arg='input',
1686
    ...                               input=im, structure=ball(3), cores=1)
1687
    >>> im3 = spim.binary_dilation(input=im, structure=ball(3))
1688
    >>> np.all(im2 == im3)
1689
    True
1690

1691
    """
1692

1693 2
    @dask.delayed
1694
    def apply_func(func, **kwargs):
1695
        # Apply function on sub-slice of overall image
1696 2
        return func(**kwargs)
1697

1698
    # Import the array_split methods
1699 2
    from array_split import shape_split, ARRAY_BOUNDS
1700

1701
    # Determine the value for im_arg
1702 2
    if type(im_arg) == str:
1703 2
        im_arg = [im_arg]
1704 2
    for item in im_arg:
1705 2
        if item in kwargs.keys():
1706 2
            im = kwargs[item]
1707 2
            im_arg = item
1708 2
            break
1709
    # Fetch image from the kwargs dict
1710 2
    im = kwargs[im_arg]
1711
    # Determine the number of divisions to create
1712 2
    divs = np.ones((im.ndim,), dtype=int) * np.array(divs)
1713
    # If overlap given then use it, otherwise search for strel in kwargs
1714 2
    if overlap is not None:
1715 2
        halo = overlap * (divs > 1)
1716
    else:
1717 2
        if type(strel_arg) == str:
1718 2
            strel_arg = [strel_arg]
1719 2
        for item in strel_arg:
1720 2
            if item in kwargs.keys():
1721 2
                strel = kwargs[item]
1722 2
                break
1723 2
        halo = np.array(strel.shape) * (divs > 1)
1724 2
    slices = np.ravel(
1725
        shape_split(
1726
            im.shape, axis=divs, halo=halo.tolist(), tile_bounds_policy=ARRAY_BOUNDS
1727
        )
1728
    )
1729
    # Apply func to each subsection of the image
1730 2
    res = []
1731
    # print('Image will be broken into the following chunks:')
1732 2
    for s in slices:
1733
        # Extract subsection from image and input into kwargs
1734 2
        kwargs[im_arg] = im[tuple(s)]
1735
        # print(kwargs[im_arg].shape)
1736 2
        res.append(apply_func(func=func, **kwargs))
1737
    # Have dask actually compute the function on each subsection in parallel
1738
    # with ProgressBar():
1739
    #    ims = dask.compute(res, num_workers=cores)[0]
1740 2
    ims = dask.compute(res, num_workers=cores)[0]
1741
    # Finally, put the pieces back together into a single master image, im2
1742 2
    im2 = np.zeros_like(im, dtype=im.dtype)
1743 2
    for i, s in enumerate(slices):
1744
        # Prepare new slice objects into main and sub-sliced image
1745 2
        a = []  # Slices into main image
1746 2
        b = []  # Slices into chunked image
1747 2
        for dim in range(im.ndim):
1748 2
            if s[dim].start == 0:
1749 2
                ax = bx = 0
1750
            else:
1751 2
                ax = s[dim].start + halo[dim]
1752 2
                bx = halo[dim]
1753 2
            if s[dim].stop == im.shape[dim]:
1754 2
                ay = by = im.shape[dim]
1755
            else:
1756 2
                ay = s[dim].stop - halo[dim]
1757 2
                by = s[dim].stop - s[dim].start - halo[dim]
1758 2
            a.append(slice(ax, ay, None))
1759 2
            b.append(slice(bx, by, None))
1760
        # Convert lists of slices to tuples
1761 2
        a = tuple(a)
1762 2
        b = tuple(b)
1763
        # Insert image chunk into main image
1764 2
        im2[a] = ims[i][b]
1765 2
    return im2
1766

1767

1768 2
def snow_partitioning_parallel(im,
1769
                               overlap='dt',
1770
                               divs=2,
1771
                               mode='parallel',
1772
                               num_workers=None,
1773
                               crop=True,
1774
                               zoom_factor=0.5,
1775
                               r_max=5,
1776
                               sigma=0.4,
1777
                               return_all=False):
1778
    r"""
1779
    Perform SNOW algorithm in parallel and serial mode to reduce time and
1780
    memory usage repectively by geomertirc domain decomposition of large size
1781
    image.
1782

1783
    Parameters
1784
    ----------
1785
    im: ND_array
1786
        A binary image of porous media with 'True' values indicating phase of
1787
        interest
1788

1789
    overlap: float or int
1790
        Overlapping thickness between two subdomains that is used to merge
1791
        watershed segmented regions at intersection of two or more subdomains.
1792
        If 'dt' the overlap will be calculated based on maximum
1793
        distance transform in the whole image.
1794
        If 'ws' the overlap will be calculated by finding the maximum dimension
1795
        of the bounding box of largest segmented region. The image is scale down
1796
        by 'zoom_factor' provided by user.
1797
        If any real number of overlap is provided then this value will be
1798
        considered as overlapping thickness.
1799

1800
    divs: list or int
1801
        Number of domains each axis will be divided.
1802
        If a scalar is provided then it will be assigned to all axis.
1803
        If list is provided then each respective axis will be divided by its
1804
        corresponding number in the list. For example [2, 3, 4] will divide
1805
        z, y and x axis to 2, 3, and 4 respectively.
1806

1807
    mode: str
1808
        if 'parallel' then all subdomains will be processed in number of cores
1809
        provided as num_workers
1810
        if 'serial' then all subdomains will be processed one by one in one core
1811
        of CPU.
1812

1813
    num_workers: int
1814
        Number of cores that will be used to parallel process all domains.
1815
        If None then all cores will be used but user can specify any integer
1816
        values to control the memory usage.
1817

1818
    crop: bool
1819
        If True the image shape is cropped to fit specified division.
1820

1821
    zoom_factor: float or int
1822
        The amount of zoom appiled to image to find overlap thickness using "ws"
1823
        overlap mode.
1824

1825
    return_all : boolean
1826
        If set to ``True`` a named tuple is returned containing the original
1827
        image, the distance transform, and the final
1828
        pore regions.  The default is ``False``
1829

1830
    Returns
1831
    ----------
1832
    regions: ND_array
1833
        Partitioned image of segmentated regions with unique labels. Each
1834
        region correspond to pore body while intersection with other region
1835
        correspond throat area.
1836
    """
1837
    # --------------------------------------------------------------------------
1838
    # Adjust image shape according to specified dimension
1839 2
    tup = namedtuple("results", field_names=["im", "dt", "regions"])
1840 2
    if isinstance(divs, int):
1841 0
        divs = [divs for i in range(im.ndim)]
1842 2
    shape = []
1843 2
    for i in range(im.ndim):
1844 2
        shape.append(divs[i] * (im.shape[i] // divs[i]))
1845

1846 2
    if shape != list(im.shape):
1847 0
        if crop:
1848 0
            for i in range(im.ndim):
1849 0
                im = im.swapaxes(0, i)
1850 0
                im = im[:shape[i], ...]
1851 0
                im = im.swapaxes(i, 0)
1852 0
            print(f'Image is cropped to shape {shape}.')
1853
        else:
1854 0
            print('-' * 80)
1855 0
            print(f"Possible image shape for specified divisions is {shape}.")
1856 0
            print("To crop the image please set crop argument to 'True'.")
1857 0
            return
1858
    # --------------------------------------------------------------------------
1859
    # Get overlap thickness from distance transform
1860 2
    chunk_shape = (np.array(shape) / np.array(divs)).astype(int)
1861 2
    print('# Beginning parallel SNOW algorithm...')
1862 2
    print('=' * 80)
1863 2
    print('Calculating overlap thickness')
1864 2
    if overlap == 'dt':
1865 2
        dt = edt((im > 0), parallel=0)
1866 2
        overlap = dt.max()
1867 0
    elif overlap == 'ws':
1868 0
        rev = spim.interpolation.zoom(im, zoom=zoom_factor, order=0)
1869 0
        rev = rev > 0
1870 0
        dt = edt(rev, parallel=0)
1871 0
        rev_snow = snow_partitioning(rev, dt=dt, r_max=r_max, sigma=sigma)
1872 0
        labels, counts = np.unique(rev_snow, return_counts=True)
1873 0
        node = np.where(counts == counts[1:].max())[0][0]
1874 0
        slices = spim.find_objects(rev_snow)
1875 0
        overlap = max(rev_snow[slices[node - 1]].shape) / (zoom_factor * 2.0)
1876 0
        dt = edt((im > 0), parallel=0)
1877
    else:
1878 0
        overlap = overlap / 2.0
1879 0
        dt = edt((im > 0), parallel=0)
1880 2
    print('Overlap Thickness: ' + str(int(2.0 * overlap)) + ' voxels')
1881
    # --------------------------------------------------------------------------
1882
    # Get overlap and trim depth of all image dimension
1883 2
    depth = {}
1884 2
    trim_depth = {}
1885 2
    for i in range(im.ndim):
1886 2
        depth[i] = int(2.0 * overlap)
1887 2
        trim_depth[i] = int(2.0 * overlap) - 1
1888

1889 2
    tup.im = im
1890 2
    tup.dt = dt
1891
    # --------------------------------------------------------------------------
1892
    # Applying snow to image chunks
1893 2
    im = da.from_array(dt, chunks=chunk_shape)
1894 2
    im = da.overlap.overlap(im, depth=depth, boundary='none')
1895 2
    im = im.map_blocks(chunked_snow, r_max=r_max, sigma=sigma)
1896 2
    im = da.overlap.trim_internal(im, trim_depth, boundary='none')
1897 2
    if mode == 'serial':
1898 0
        num_workers = 1
1899 2
    elif mode == 'parallel':
1900 2
        num_workers = num_workers
1901
    else:
1902 0
        raise Exception('Mode of operation can either be parallel or serial')
1903 2
    with ProgressBar():
1904
        # print('-' * 80)
1905 2
        print('Applying snow to image chunks')
1906 2
        regions = im.compute(num_workers=num_workers)
1907
    # --------------------------------------------------------------------------
1908
    # Relabelling watershed chunks
1909
    # print('-' * 80)
1910 2
    print('Relabelling watershed chunks')
1911 2
    regions = relabel_chunks(im=regions, chunk_shape=chunk_shape)
1912
    # --------------------------------------------------------------------------
1913
    # Stitching watershed chunks
1914
    # print('-' * 80)
1915 2
    print('Stitching watershed chunks')
1916 2
    regions = watershed_stitching(im=regions, chunk_shape=chunk_shape)
1917 2
    print('=' * 80)
1918 2
    if return_all:
1919 2
        tup.regions = regions
1920 2
        return tup
1921
    else:
1922 0
        return regions
1923

1924 0
    return regions
1925

1926

1927 2
def chunked_snow(im, r_max=5, sigma=0.4):
1928
    r"""
1929
    Partitions the void space into pore regions using a marker-based watershed
1930
    algorithm, with specially filtered peaks as markers.
1931

1932
    The SNOW network extraction algorithm (Sub-Network of an Over-segmented
1933
    Watershed) was designed to handle to perculiarities of high porosity
1934
    materials, but it applies well to other materials as well.
1935

1936
    Parameters
1937
    ----------
1938
    im : array_like
1939
        Distance transform of phase of interest in a binary image
1940
    r_max : int
1941
        The radius of the spherical structuring element to use in the Maximum
1942
        filter stage that is used to find peaks.  The default is 5
1943
    sigma : float
1944
        The standard deviation of the Gaussian filter used in step 1.  The
1945
        default is 0.4.  If 0 is given then the filter is not applied, which is
1946
        useful if a distance transform is supplied as the ``im`` argument that
1947
        has already been processed.
1948

1949
    Returns
1950
    -------
1951
    image : ND-array
1952
        An image the same shape as ``im`` with the void space partitioned into
1953
        pores using a marker based watershed with the peaks found by the
1954
        SNOW algorithm [1].
1955

1956
    References
1957
    ----------
1958
    [1] Gostick, J. "A versatile and efficient network extraction algorithm
1959
    using marker-based watershed segmenation".  Physical Review E. (2017)
1960
    """
1961

1962 2
    dt = spim.gaussian_filter(input=im, sigma=sigma)
1963 2
    peaks = find_peaks(dt=dt, r_max=r_max)
1964 2
    peaks = trim_saddle_points(peaks=peaks, dt=dt, max_iters=99, verbose=0)
1965 2
    peaks = trim_nearby_peaks(peaks=peaks, dt=dt)
1966 2
    peaks, N = spim.label(peaks)
1967 2
    regions = watershed(image=-dt, markers=peaks, mask=im > 0)
1968

1969 2
    return regions * (im > 0)
1970

1971

1972 2
def pad(im, pad_width=1, constant_value=0):
1973
    r"""
1974
    Pad the image with a constant values and width.
1975

1976
    Parameters:
1977
    ----------
1978
    im : ND-array
1979
        The image that requires padding
1980
    pad_width : int
1981
        The number of values that will be padded from the edges. Default values
1982
        is 1.
1983
    contant_value : int
1984
        Pads with the specified constant value
1985

1986
    return:
1987
    -------
1988
    output: ND-array
1989
        Padded image with same dimnesions as provided image
1990
    """
1991 2
    shape = np.array(im.shape)
1992 2
    pad_shape = shape + (2 * pad_width)
1993 2
    temp = np.zeros(pad_shape, dtype=np.uint32)
1994 2
    if constant_value != 0:
1995 0
        temp = temp + constant_value
1996 2
    if im.ndim == 3:
1997 0
        temp[pad_width: -pad_width,
1998
             pad_width: -pad_width,
1999
             pad_width: -pad_width] = im
2000 2
    elif im.ndim == 2:
2001 2
        temp[pad_width: -pad_width,
2002
             pad_width: -pad_width] = im
2003
    else:
2004 0
        temp[pad_width: -pad_width] = im
2005

2006 2
    return temp
2007

2008

2009 2
def relabel_chunks(im, chunk_shape):
2010
    r"""
2011
    Assign new labels to each chunk or sub-domain of actual image. This
2012
    prevents from two or more regions to have same label.
2013

2014
    Parameters:
2015
    -----------
2016

2017
    im: ND-array
2018
        Actual image that contains repeating labels in chunks or sub-domains
2019

2020
    chunk_shape: tuple
2021
        The shape of chunk that will be relabeled in actual image. Note the
2022
        chunk shape should be a multiple of actual image shape otherwise some
2023
        labels will not be relabeled.
2024

2025
    return:
2026
    -------
2027
    output : ND-array
2028
        Relabeled image with unique label assigned to each region.
2029
    """
2030 2
    im = pad(im, pad_width=1)
2031 2
    im_shape = np.array(im.shape, dtype=np.uint32)
2032 2
    max_num = 0
2033 2
    c = np.array(chunk_shape, dtype=np.uint32) + 2
2034 2
    num = (im_shape / c).astype(int)
2035

2036 2
    if im.ndim == 3:
2037 0
        for z in range(num[0]):
2038 0
            for y in range(num[1]):
2039 0
                for x in range(num[2]):
2040 0
                    chunk = im[z * c[0]: (z + 1) * c[0],
2041
                               y * c[1]: (y + 1) * c[1],
2042
                               x * c[2]: (x + 1) * c[2]]
2043 0
                    chunk += max_num
2044 0
                    chunk[chunk == max_num] = 0
2045 0
                    max_num = chunk.max()
2046 0
                    im[z * c[0]: (z + 1) * c[0],
2047
                       y * c[1]: (y + 1) * c[1],
2048
                       x * c[2]: (x + 1) * c[2]] = chunk
2049
    else:
2050 2
        for y in range(num[0]):
2051 2
            for x in range(num[1]):
2052 2
                chunk = im[y * c[0]: (y + 1) * c[0],
2053
                           x * c[1]: (x + 1) * c[1]]
2054 2
                chunk += max_num
2055 2
                chunk[chunk == max_num] = 0
2056 2
                max_num = chunk.max()
2057 2
                im[y * c[0]: (y + 1) * c[0],
2058
                   x * c[1]: (x + 1) * c[1]] = chunk
2059

2060 2
    return im
2061

2062

2063 2
def trim_internal_slice(im, chunk_shape):
2064
    r"""
2065
    Delete extra slices from image that were used to stitch two or more chunks
2066
    together.
2067

2068
    Parameters:
2069
    -----------
2070

2071
    im :  ND-array
2072
        image that contains extra slices in x, y, z direction.
2073

2074
    chunk_shape : tuple
2075
        The shape of the chunk from which image is subdivided.
2076

2077
    Return:
2078
    -------
2079
    output :  ND-array
2080
        Image without extra internal slices. The shape of the image will be
2081
        same as input image provided for  waterhsed segmentation.
2082
    """
2083 2
    im_shape = np.array(im.shape, dtype=np.uint32)
2084 2
    c1 = np.array(chunk_shape, dtype=np.uint32) + 2
2085 2
    c2 = np.array(chunk_shape, dtype=np.uint32)
2086 2
    num = (im_shape / c1).astype(int)
2087 2
    out_shape = num * c2
2088 2
    out = np.empty((out_shape), dtype=np.uint32)
2089

2090 2
    if im.ndim == 3:
2091 0
        for z in range(num[0]):
2092 0
            for y in range(num[1]):
2093 0
                for x in range(num[2]):
2094 0
                    chunk = im[z * c1[0]: (z + 1) * c1[0],
2095
                               y * c1[1]: (y + 1) * c1[1],
2096
                               x * c1[2]: (x + 1) * c1[2]]
2097

2098 0
                    out[z * c2[0]: (z + 1) * c2[0],
2099
                        y * c2[1]: (y + 1) * c2[1],
2100
                        x * c2[2]: (x + 1) * c2[2]] = chunk[1:-1, 1:-1, 1:-1]
2101
    else:
2102 2
        for y in range(num[0]):
2103 2
            for x in range(num[1]):
2104 2
                chunk = im[y * c1[0]: (y + 1) * c1[0],
2105
                           x * c1[1]: (x + 1) * c1[1]]
2106

2107 2
                out[y * c2[0]: (y + 1) * c2[0],
2108
                    x * c2[1]: (x + 1) * c2[1]] = chunk[1:-1, 1:-1]
2109

2110 2
    return out
2111

2112

2113 2
def watershed_stitching(im, chunk_shape):
2114
    r"""
2115
    Stitch individual sub-domains of watershed segmentation into one big
2116
    segmentation with all boundary labels of each sub-domain relabeled to merge
2117
    boundary regions.
2118

2119
    Parameters:
2120
    -----------
2121
    im : ND-array
2122
        A worked image with watershed segmentation performed on all sub-domains
2123
        individually.
2124

2125
    chunk_shape: tuple
2126
        The shape of the sub-domain in which image segmentation is performed.
2127

2128
    return:
2129
    -------
2130
    output : ND-array
2131
        Stitched watershed segmentation with all sub-domains merged to form a
2132
        single watershed segmentation.
2133
    """
2134

2135 2
    c_shape = np.array(chunk_shape)
2136 2
    cuts_num = (np.array(im.shape) / c_shape).astype(np.uint32)
2137

2138 2
    for axis, num in enumerate(cuts_num):
2139 2
        keys = []
2140 2
        values = []
2141 2
        if num > 1:
2142 2
            im = im.swapaxes(0, axis)
2143 2
            for i in range(1, num):
2144 2
                sl = i * (chunk_shape[axis] + 3) - (i - 1)
2145 2
                sl1 = im[sl - 3, ...]
2146 2
                sl1_mask = sl1 > 0
2147 2
                sl2 = im[sl - 1, ...] * sl1_mask
2148 2
                sl1_labels = sl1.flatten()[sl1.flatten() > 0]
2149 2
                sl2_labels = sl2.flatten()[sl2.flatten() > 0]
2150 2
                if sl1_labels.size != sl2_labels.size:
2151 0
                    raise Exception('The selected overlapping thickness is not '
2152
                                    'suitable for input image. Change '
2153
                                    'overlapping criteria '
2154
                                    'or manually input value.')
2155 2
                keys.append(sl1_labels)
2156 2
                values.append(sl2_labels)
2157 2
            im = replace_labels(array=im, keys=keys, values=values)
2158 2
            im = im.swapaxes(axis, 0)
2159 2
    im = trim_internal_slice(im=im, chunk_shape=chunk_shape)
2160 2
    im = resequence_labels(array=im)
2161

2162 2
    return im
2163

2164

2165 2
@njit(parallel=True)
2166
def copy(im, output):  # pragma: no cover
2167
    r"""
2168
    The function copy the input array and make output array that is allocated
2169
    in different memory space. This a numba version of copy function of numpy.
2170
    Because each element is copied using parallel approach this implementation
2171
    is facter than numpy version of copy.
2172

2173
    parameter:
2174
    ----------
2175
    array: ND-array
2176
        Array that needs to be copied
2177

2178
    Return:
2179
    -------
2180
    output: ND-array
2181
        Copied array
2182
    """
2183

2184
    if im.ndim == 3:
2185
        for i in prange(im.shape[0]):
2186
            for j in prange(im.shape[1]):
2187
                for k in prange(im.shape[2]):
2188
                    output[i, j, k] = im[i, j, k]
2189
    elif im.ndim == 2:
2190
        for i in prange(im.shape[0]):
2191
            for j in prange(im.shape[1]):
2192
                output[i, j] = im[i, j]
2193
    else:
2194
        for i in prange(im.shape[0]):
2195
            output[i] = im[i]
2196

2197
    return output
2198

2199

2200 2
@njit(parallel=True)
2201
def _replace(array, keys, values, ind_sort):  # pragma: no cover
2202
    r"""
2203
    This function replace keys elements in input array with new value elements.
2204
    This function is used as internal function of replace_relabels.
2205

2206
    Parameter:
2207
    ----------
2208
    array : ND-array
2209
        Array which requires replacing labels
2210
    keys :  1D-array
2211
        The unique labels that need to be replaced
2212
    values : 1D-array
2213
        The unique values that will be assigned to labels
2214

2215
    return:
2216
    -------
2217
    array : ND-array
2218
        Array with replaced labels.
2219
    """
2220
    # ind_sort = np.argsort(keys)
2221
    keys_sorted = keys[ind_sort]
2222
    values_sorted = values[ind_sort]
2223
    s_keys = set(keys)
2224

2225
    for i in prange(array.shape[0]):
2226
        if array[i] in s_keys:
2227
            ind = np.searchsorted(keys_sorted, array[i])
2228
            array[i] = values_sorted[ind]
2229

2230

2231 2
def replace_labels(array, keys, values):
2232
    r"""
2233
    Replace labels in array provided as keys to values.
2234

2235
    Parameter:
2236
    ----------
2237
    array : ND-array
2238
        Array which requires replacing labels
2239
    keys :  1D-array
2240
        The unique labels that need to be replaced
2241
    values : 1D-array
2242
        The unique values that will be assigned to labels
2243

2244
    return:
2245
    -------
2246
    array : ND-array
2247
        Array with replaced labels.
2248
    """
2249 2
    a_shape = array.shape
2250 2
    array = array.flatten()
2251 2
    keys = np.concatenate(keys, axis=0)
2252 2
    values = np.concatenate(values, axis=0)
2253 2
    ind_sort = np.argsort(keys)
2254 2
    _replace(array, keys, values, ind_sort)
2255

2256 2
    return array.reshape(a_shape)
2257

2258

2259 2
@njit()
2260
def _sequence(array, count):  # pragma: no cover
2261
    r"""
2262
    Internal function of resequnce_labels method. This function resquence array
2263
    elements in an ascending order using numba technique which is many folds
2264
    faster than make contigious funcition.
2265

2266
    parameter:
2267
    ----------
2268
    array: 1d-array
2269
        1d-array that needs resquencing
2270
    count: 1d-array
2271
        1d-array of zeros having same size as array
2272

2273
    return:
2274
    -------
2275
    array: 1d-array
2276
        The input array with elements resequenced in ascending order
2277
    Note: The output of this function is not same as make_contigous or
2278
    relabel_sequential function of scikit-image. This function resequence and
2279
    randomize the regions while other methods only do resequencing and output
2280
    sorted array.
2281
    """
2282
    a = 1
2283
    i = 0
2284
    while i < (len(array)):
2285
        data = array[i]
2286
        if data != 0:
2287
            if count[data] == 0:
2288
                count[data] = a
2289
                a += 1
2290
        array[i] = count[data]
2291
        i += 1
2292

2293

2294 2
@njit(parallel=True)
2295
def amax(array):  # pragma: no cover
2296
    r"""
2297
    Find largest element in an array using fast parallel numba technique
2298

2299
    Parameter:
2300
    ----------
2301
    array: ND-array
2302
        array in which largest elements needs to be calcuted
2303

2304
    return:
2305
    scalar: float or int
2306
        The largest element value in the input array
2307
    """
2308

2309
    return np.max(array)
2310

2311

2312 2
def resequence_labels(array):
2313
    r"""
2314
    Resequence the lablels to make them contigious.
2315

2316
    Parameter:
2317
    ----------
2318
    array: ND-array
2319
        Array that requires resequencing
2320

2321
    return:
2322
    -------
2323
    array : ND-array
2324
        Resequenced array with same shape as input array
2325
    """
2326 2
    a_shape = array.shape
2327 2
    array = array.ravel()
2328 2
    max_num = amax(array) + 1
2329 2
    count = np.zeros(max_num, dtype=np.uint32)
2330 2
    _sequence(array, count)
2331

2332 2
    return array.reshape(a_shape)

Read our documentation on viewing source code .

Loading