numba-morph 0.1.0__py3-none-any.whl

numba_morph/__init__.py

@@ -0,0 +1,15 @@
+__version__ = "0.1.0"
+
+from .cdt import distance_transform_cdt
+from .dilation import dilation
+from .erosion import erosion
+from .gradient import morphological_gradient
+from .h_min_max import h_minima, h_maxima
+from .laplace import morphological_laplace
+from .local_min_max import local_minima, local_maxima
+from .open_close import opening, closing
+from .reconstruction import reconstruction
+from .top_hats import white_tophat, black_tophat
+from .utils import generate_sphere_structure
+from .watershed import watershed
+from .welford import welford_mean_std_w_mask

numba_morph/_chamfer.py

@@ -0,0 +1,409 @@
+import numpy as np
+from numba import njit, prange
+
+
+def _generate_offsets(weights):
+    """
+    Generate causal and anti-causal offsets for Chamfer masks.
+    weights : tuple of ints
+        For 2D: (face_weight, edge_weight)
+        For 3D: (face_weight, edge_weight, corner_weight)
+    Returns two arrays of shape (N, ndim+1) with columns [offset_0, ..., offset_{ndim-1}, weight].
+    """
+    ndim = len(weights)
+    causal = []
+    anti_causal = []
+
+    if ndim == 2:
+        # 2D offsets: (dy, dx)
+        for dy in (-1, 0, 1):
+            for dx in (-1, 0, 1):
+                if dy == 0 and dx == 0:
+                    continue
+                manhattan = abs(dy) + abs(dx)
+                if manhattan > ndim:          # max Manhattan distance is 2
+                    continue
+                w = weights[manhattan - 1]
+                # causal: first non‑zero component is negative
+                if dy < 0 or (dy == 0 and dx < 0):
+                    causal.append((dy, dx, w))
+                if dy > 0 or (dy == 0 and dx > 0):
+                    anti_causal.append((dy, dx, w))
+
+    elif ndim == 3:
+        # 3D offsets: (dz, dy, dx)
+        for dz in (-1, 0, 1):
+            for dy in (-1, 0, 1):
+                for dx in (-1, 0, 1):
+                    if dz == 0 and dy == 0 and dx == 0:
+                        continue
+                    manhattan = abs(dz) + abs(dy) + abs(dx)
+                    if manhattan > ndim:       # max Manhattan distance is 3
+                        continue
+                    w = weights[manhattan - 1]
+                    if dz < 0 or (dz == 0 and dy < 0) or (dz == 0 and dy == 0 and dx < 0):
+                        causal.append((dz, dy, dx, w))
+                    if dz > 0 or (dz == 0 and dy > 0) or (dz == 0 and dy == 0 and dx > 0):
+                        anti_causal.append((dz, dy, dx, w))
+    else:
+        raise ValueError("Only 2D or 3D are supported")
+
+    return np.array(causal, dtype=np.int32), np.array(anti_causal, dtype=np.int32)
+
+
+@njit(cache=True)
+def _apply_pass_2d(input, offsets,
+                   y_start, y_end, y_step,
+                   x_start, x_end, x_step,
+                   H, W, max_val):
+    changed = False
+    for y in range(y_start, y_end, y_step):
+        for x in range(x_start, x_end, x_step):
+            if input[y, x] == 0:
+                continue
+            cur = input[y, x]
+            new_val = cur
+            for dy, dx, w in offsets:
+                ny = y + dy
+                nx = x + dx
+                if 0 <= ny < H and 0 <= nx < W:
+                    neigh = input[ny, nx]
+                    if neigh != max_val:
+                        cand = neigh + w
+                        if cand < new_val:
+                            new_val = cand
+            if new_val < cur:
+                input[y, x] = new_val
+                changed = True
+    return changed
+
+
+@njit(parallel=True, cache=True)
+def _apply_pass_2d_batch(input, offsets,
+                         y_start, y_end, y_step,
+                         x_start, x_end, x_step,
+                         L, H, W, max_val):
+    changed = 0
+    for l in prange(L):
+        for y in range(y_start, y_end, y_step):
+            for x in range(x_start, x_end, x_step):
+                if input[l, y, x] == 0:
+                    continue
+                cur = input[l, y, x]
+                new_val = cur
+                for dy, dx, w in offsets:
+                    ny = y + dy
+                    nx = x + dx
+                    if 0 <= ny < H and 0 <= nx < W:
+                        neigh = input[l, ny, nx]
+                        if neigh != max_val:
+                            cand = neigh + w
+                            if cand < new_val:
+                                new_val = cand
+                if new_val < cur:
+                    input[l, y, x] = new_val
+                    changed += 1
+    return True if changed > 0 else False
+
+
+@njit(cache=True)
+def _apply_pass_3d(input, offsets,
+                   z_start, z_end, z_step,
+                   y_start, y_end, y_step,
+                   x_start, x_end, x_step,
+                   D, H, W, max_val):
+    changed = False
+    for z in range(z_start, z_end, z_step):
+        for y in range(y_start, y_end, y_step):
+            for x in range(x_start, x_end, x_step):
+                if input[z, y, x] == 0:
+                    continue
+                cur = input[z, y, x]
+                new_val = cur
+                for dz, dy, dx, w in offsets:
+                    nz = z + dz
+                    ny = y + dy
+                    nx = x + dx
+                    if 0 <= nz < D and 0 <= ny < H and 0 <= nx < W:
+                        neigh = input[nz, ny, nx]
+                        if neigh != max_val:
+                            cand = neigh + w
+                            if cand < new_val:
+                                new_val = cand
+                if new_val < cur:
+                    input[z, y, x] = new_val
+                    changed = True
+    return changed
+
+
+@njit(parallel=True, cache=True)
+def _apply_pass_3d_batch(input, offsets,
+                         z_start, z_end, z_step,
+                         y_start, y_end, y_step,
+                         x_start, x_end, x_step,
+                         L, D, H, W, max_val):
+    changed = 0
+    for l in prange(L):
+        for z in range(z_start, z_end, z_step):
+            for y in range(y_start, y_end, y_step):
+                for x in range(x_start, x_end, x_step):
+                    if input[l, z, y, x] == 0:
+                        continue
+                    cur = input[l, z, y, x]
+                    new_val = cur
+                    for dz, dy, dx, w in offsets:
+                        nz = z + dz
+                        ny = y + dy
+                        nx = x + dx
+                        if 0 <= nz < D and 0 <= ny < H and 0 <= nx < W:
+                            neigh = input[l, nz, ny, nx]
+                            if neigh != max_val:
+                                cand = neigh + w
+                                if cand < new_val:
+                                    new_val = cand
+                    if new_val < cur:
+                        input[l, z, y, x] = new_val
+                        changed += 1
+    return True if changed > 0 else False
+
+
+@njit(parallel=True, fastmath=True, cache=True)
+def _chamfer_2d_chunk(input, max_val, num_bands, causal, anti_causal, size_of_largest_dim):
+    H, W = input.shape
+    chunk_size = (size_of_largest_dim // num_bands) + 1
+
+    for it in range(size_of_largest_dim):
+        changed = np.zeros(num_bands, dtype=np.bool_)
+
+        for band in prange(num_bands):
+            start = band * chunk_size
+            end = min(start + chunk_size, size_of_largest_dim)
+
+            # forward pass
+            if _apply_pass_2d(input, causal,
+                              start, end, 1,  # y: start -> end-1
+                              0, W, 1,  # x: 0 -> W-1
+                              H, W, max_val):
+                changed[band] = True
+
+            # backward pass
+            if _apply_pass_2d(input, anti_causal,
+                             end - 1, start - 1, -1,  # y: end-1 down to start
+                             W - 1, -1, -1,  # x: W-1 down to 0
+                              H, W, max_val):
+                changed[band] = True
+
+        if not np.any(changed):
+            break
+
+
+@njit(parallel=True, fastmath=True, cache=True)
+def _chamfer_2d_chunk_batch(input, max_val, num_bands, causal, anti_causal, size_of_largest_dim):
+    L, H, W = input.shape
+    chunk_size = (size_of_largest_dim // num_bands) + 1
+
+    for it in range(size_of_largest_dim):
+        changed = np.zeros(num_bands, dtype=np.bool_)
+
+        for band in prange(num_bands):
+            start = band * chunk_size
+            end = min(start + chunk_size, size_of_largest_dim)
+
+            # forward pass
+            if _apply_pass_2d_batch(input, causal,
+                                    start, end, 1,  # y: start -> end-1
+                                    0, W, 1,  # x: 0 -> W-1
+                                    L, H, W, max_val):
+                changed[band] = True
+
+            # backward pass
+            if _apply_pass_2d_batch(input, anti_causal,
+                                    end - 1, start - 1, -1,  # y: end-1 down to start
+                                    W - 1, -1, -1,  # x: W-1 down to 0
+                                    L, H, W, max_val):
+                changed[band] = True
+
+        if not np.any(changed):
+            break
+
+
+@njit(parallel=True, fastmath=True, cache=True)
+def _chamfer_3d_chunk(input, max_val, num_bands, causal, anti_causal, size_of_largest_dim):
+
+    D, H, W = input.shape
+    chunk_size = (size_of_largest_dim // num_bands) + 1
+
+    for it in range(size_of_largest_dim):
+        changed = np.zeros(num_bands, dtype=np.bool_)
+
+        for band in prange(num_bands):
+            start = band * chunk_size
+            end = min(start + chunk_size, size_of_largest_dim)
+
+            # forward pass
+            if _apply_pass_3d(input, causal,
+                              start, end, 1,  # y: start -> end-1
+                              0, H, 1,
+                              0, W, 1,  # x: 0 -> W-1
+                              D, H, W, max_val):
+                changed[band] = True
+
+            # backward pass
+            if _apply_pass_3d(input, anti_causal,
+                              end - 1, start - 1, -1,  # y: end-1 down to start
+                              H - 1, -1, -1,
+                              W - 1, -1, -1,  # x: W-1 down to 0
+                              D, H, W, max_val):
+                changed[band] = True
+
+        if not np.any(changed):
+            break
+
+
+@njit(parallel=True, fastmath=True, cache=True)
+def _chamfer_3d_chunk_batch(input, max_val, num_bands, causal, anti_causal, size_of_largest_dim):
+    L, D, H, W = input.shape
+    chunk_size = (size_of_largest_dim // num_bands) + 1
+
+    for it in range(size_of_largest_dim):
+        changed = np.zeros(num_bands, dtype=np.bool_)
+
+        for band in prange(num_bands):
+            start = band * chunk_size
+            end = min(start + chunk_size, size_of_largest_dim)
+
+            # forward pass
+            if _apply_pass_3d_batch(input, causal,
+                                    start, end, 1,
+                                    0, H, 1,
+                                    0, W, 1,  # x: 0 -> W-1
+                                    L, D, H, W, max_val):
+                changed[band] = True
+
+            # backward pass
+            if _apply_pass_3d_batch(input, anti_causal,
+                                    end - 1, start - 1, -1,
+                                    H - 1, -1, -1,
+                                    W - 1, -1, -1,  # x: W-1 down to 0
+                                    L, D, H, W, max_val):
+                changed[band] = True
+
+        if not np.any(changed):
+            break
+
+
+def reorder_offset_list(offset_list, spatial_perm):
+    new_list = []
+    for item in offset_list:
+        spatial_offsets = item[:-1]  # all but the last element (weight)
+        weight = item[-1]
+        # Reorder spatial offsets according to spatial_perm
+        new_spatial = tuple(spatial_offsets[i] for i in spatial_perm)
+        new_list.append((*new_spatial, weight))
+    return new_list
+
+
+def _make_ranges(shape):
+    """Return (fwd_ranges, bwd_ranges) for a given spatial shape."""
+    fwd = tuple((0, d, 1) for d in shape)
+    bwd = tuple((d - 1, -1, -1) for d in shape)
+    return fwd, bwd
+
+
+# ---------- Dispatcher for a single forward/backward pass ----------
+def _run_pass(input, causal, anti_causal, ranges_fwd, ranges_bwd, ndim, batch, max_val):
+    """
+    Call the appropriate _apply_pass_* function with the given ranges.
+    'ranges_fwd' and 'ranges_bwd' are tuples of (start, end, step) per axis.
+    """
+    if ndim == 2:
+        if batch:
+            func = _apply_pass_2d_batch
+            # ranges: (y_range, x_range)
+            (y_s, y_e, y_st), (x_s, x_e, x_st) = ranges_fwd
+            func(input, causal, y_s, y_e, y_st, x_s, x_e, x_st, input.shape[0], input.shape[1], input.shape[2], max_val)
+            (y_s, y_e, y_st), (x_s, x_e, x_st) = ranges_bwd
+            func(input, anti_causal, y_s, y_e, y_st, x_s, x_e, x_st, input.shape[0], input.shape[1], input.shape[2], max_val)
+        else:
+            func = _apply_pass_2d
+            (y_s, y_e, y_st), (x_s, x_e, x_st) = ranges_fwd
+            func(input, causal, y_s, y_e, y_st, x_s, x_e, x_st, input.shape[0], input.shape[1], max_val)
+            (y_s, y_e, y_st), (x_s, x_e, x_st) = ranges_bwd
+            func(input, anti_causal, y_s, y_e, y_st, x_s, x_e, x_st, input.shape[0], input.shape[1], max_val)
+    else:  # ndim == 3
+        if batch:
+            func = _apply_pass_3d_batch
+            (z_s, z_e, z_st), (y_s, y_e, y_st), (x_s, x_e, x_st) = ranges_fwd
+            func(input, causal, z_s, z_e, z_st, y_s, y_e, y_st, x_s, x_e, x_st,
+                 input.shape[0], input.shape[1], input.shape[2], input.shape[3], max_val)
+            (z_s, z_e, z_st), (y_s, y_e, y_st), (x_s, x_e, x_st) = ranges_bwd
+            func(input, anti_causal, z_s, z_e, z_st, y_s, y_e, y_st, x_s, x_e, x_st,
+                 input.shape[0], input.shape[1], input.shape[2], input.shape[3], max_val)
+        else:
+            func = _apply_pass_3d
+            (z_s, z_e, z_st), (y_s, y_e, y_st), (x_s, x_e, x_st) = ranges_fwd
+            func(input, causal, z_s, z_e, z_st, y_s, y_e, y_st, x_s, x_e, x_st,
+                 input.shape[0], input.shape[1], input.shape[2], max_val)
+            (z_s, z_e, z_st), (y_s, y_e, y_st), (x_s, x_e, x_st) = ranges_bwd
+            func(input, anti_causal, z_s, z_e, z_st, y_s, y_e, y_st, x_s, x_e, x_st,
+                 input.shape[0], input.shape[1], input.shape[2], max_val)
+
+
+# ---------- Dispatcher for the chunked iterative version ----------
+def _run_chunked(input, max_val, num_bands, causal, anti_causal, ndim, batch, size_of_largest_dim):
+    """Call the appropriate _chamfer_*_chunk* function."""
+    if ndim == 2:
+        if batch:
+            _chamfer_2d_chunk_batch(input, max_val, num_bands, causal, anti_causal, size_of_largest_dim)
+        else:
+            _chamfer_2d_chunk(input, max_val, num_bands, causal, anti_causal, size_of_largest_dim)
+    else:  # ndim == 3
+        if batch:
+            _chamfer_3d_chunk_batch(input, max_val, num_bands, causal, anti_causal, size_of_largest_dim)
+        else:
+            _chamfer_3d_chunk(input, max_val, num_bands, causal, anti_causal, size_of_largest_dim)
+
+
+# ---------- Helper: permute for chunking ----------
+def _permute_for_chunk(input, causal, anti_causal, working_dim, chunk_dim, batch):
+    """
+    Transpose the input so that the chunked axis becomes the first spatial axis,
+    and reorder the offsets accordingly. Returns (transposed_input, new_causal, new_anti_causal, perm, inv_perm).
+    """
+    spatial_index = working_dim + chunk_dim          # absolute axis (0-based among spatial)
+    if batch:
+        total_axes = working_dim + 1
+        abs_chunk_axis = spatial_index + 1           # batch is axis 0
+        perm = [0, abs_chunk_axis] + [i for i in range(1, total_axes) if i != abs_chunk_axis]
+        spatial_perm = [i - 1 for i in perm[1:]]     # spatial order after permutation
+    else:
+        perm = [spatial_index] + [i for i in range(working_dim) if i != spatial_index]
+        spatial_perm = perm
+
+    new_causal = reorder_offset_list(causal, spatial_perm)
+    new_anti_causal = reorder_offset_list(anti_causal, spatial_perm)
+    return np.transpose(input, perm), new_causal, new_anti_causal, perm, np.argsort(perm)
+
+
+# ---------- Refactored _chamfer ----------
+def _chamfer(input, max_val, num_bands, causal, anti_causal, working_dim,
+             chunk, size_of_largest_dim, chunk_dim, batch):
+    """
+    Compute Chamfer distance transform.
+    """
+    if chunk:
+        # Move the chunked axis to the front (after batch if present)
+        input, causal, anti_causal, perm, inv_perm = _permute_for_chunk(
+            input, causal, anti_causal, working_dim, chunk_dim, batch
+        )
+        # Run the chunked algorithm (the chunk function assumes the first spatial axis is the chunked one)
+        _run_chunked(input, max_val, num_bands, causal, anti_causal, working_dim, batch, size_of_largest_dim)
+        # Restore original axis order
+        input = np.transpose(input, inv_perm)
+    else:
+        # No chunking: simply run one forward and one backward pass over all axes
+        shape = input.shape[1:] if batch else input.shape   # spatial shape
+        fwd_ranges, bwd_ranges = _make_ranges(shape)
+        _run_pass(input, causal, anti_causal, fwd_ranges, bwd_ranges, working_dim, batch, max_val)
+
+    return input

numba_morph/_edt.py

@@ -0,0 +1,156 @@
+import numpy as np
+from numba import njit
+
+@njit
+def _voronoi_1d(site_pos, site_R, L):
+    # Collect valid sites (site_pos != -1)
+    # We'll store their indices in arrays f and g (stack)
+    f = np.empty(L, dtype=np.int32)   # site index (position in array)
+    n = 0
+    for i in range(L):
+        if site_pos[i] != -1:
+            f[n] = i
+            n += 1
+
+    best_idx = np.full(L, -1, dtype=np.int32)
+    if n == 0:
+        return best_idx
+
+    # Build the lower envelope using a stack
+    stack = np.empty(n, dtype=np.int32)
+    top = -1
+    for i in range(n):
+        # While stack has at least 2 sites, check if the last site is obsolete
+        while top >= 1:
+            # indices of the three sites: prev, last, new
+            idx1 = f[stack[top - 1]]
+            idx2 = f[stack[top]]
+            idx3 = f[i]
+
+            s1 = site_pos[idx1]
+            s2 = site_pos[idx2]
+            s3 = site_pos[idx3]
+            R1 = site_R[idx1]
+            R2 = site_R[idx2]
+            R3 = site_R[idx3]
+
+            # Intersection x between site2 and site3: (x - s2)^2 + R2 = (x - s3)^2 + R3
+            # => x23 = (R3 - R2 + s3^2 - s2^2) / (2*(s3 - s2))
+            x23 = (R3 - R2 + s3 * s3 - s2 * s2) / (2.0 * (s3 - s2))
+            # Intersection between site1 and site2
+            x12 = (R2 - R1 + s2 * s2 - s1 * s1) / (2.0 * (s2 - s1))
+
+            if x23 <= x12:
+                # last site (idx2) is never the best, pop it
+                top -= 1
+            else:
+                break
+        # Push new site
+        top += 1
+        stack[top] = i
+
+    # Now the stack contains the sites forming the lower envelope.
+    # Traverse the line and assign the best site.
+    l = 0  # pointer into stack
+    for x in range(L):
+        # Move l forward while the next site gives a smaller distance
+        while l < top:
+            idx_cur = f[stack[l]]
+            idx_next = f[stack[l + 1]]
+            s_cur = site_pos[idx_cur]
+            s_next = site_pos[idx_next]
+            R_cur = site_R[idx_cur]
+            R_next = site_R[idx_next]
+
+            # distance to current site: (x - s_cur)^2 + R_cur
+            # distance to next site: (x - s_next)^2 + R_next
+            # If next is better (or equal), advance l
+            if (x - s_next) ** 2 + R_next <= (x - s_cur) ** 2 + R_cur:
+                l += 1
+            else:
+                break
+        best_idx[x] = f[stack[l]]
+
+    return best_idx
+
+
+@njit
+def edt_feature_transform_2d(binary):
+    """
+    Compute the exact Euclidean Feature Transform for a 2D binary image.
+
+    Parameters
+    ----------
+    binary : 2D int array, shape (H, W)
+        1 for foreground, 0 for background.
+
+    Returns
+    -------
+    feat_r : 2D int array, shape (H, W)
+        Row coordinate of the nearest background pixel for each pixel.
+    feat_c : 2D int array, shape (H, W)
+        Column coordinate of the nearest background pixel for each pixel.
+        For background pixels, feat_r = row, feat_c = col (self).
+    """
+    H, W = binary.shape
+    feat_r = np.full((H, W), -1, dtype=np.int32)
+    feat_c = np.full((H, W), -1, dtype=np.int32)
+
+    # Initialise: background pixels point to themselves, foreground to -1
+    for i in range(H):
+        for j in range(W):
+            if binary[i, j] == 0:
+                feat_r[i, j] = i
+                feat_c[i, j] = j
+
+    # ----- First pass: along rows (axis 0) -----
+    # For each column, run 1D Voronoi on the rows.
+    # At this stage, accumulated R = 0.
+    for j in range(W):
+        site_pos = feat_r[:, j]      # row coordinate of site (or -1)
+        site_R = np.zeros(H, dtype=np.float64)  # no previous dims
+        best_idx = _voronoi_1d(site_pos, site_R, H)
+        for i in range(H):
+            if best_idx[i] != -1:
+                # best_idx[i] is the row index of the best site in this column
+                # The site's coordinates are (feat_r[best_idx[i], j], j)
+                feat_r[i, j] = feat_r[best_idx[i], j]
+                feat_c[i, j] = j   # column is unchanged in this pass
+
+    # ----- Second pass: along columns (axis 1) -----
+    # For each row, run 1D Voronoi on the columns.
+    # Now the accumulated R for each site is (row - feat_r[row, col])^2.
+    for i in range(H):
+        site_pos = feat_c[i, :]      # column coordinate of site (or -1)
+        # compute R = (i - feat_r[i, j])^2 for each site
+        site_R = np.zeros(W, dtype=np.float64)
+        for j in range(W):
+            if feat_r[i, j] != -1:
+                dr = i - feat_r[i, j]
+                site_R[j] = dr * dr
+        best_idx = _voronoi_1d(site_pos, site_R, W)
+        for j in range(W):
+            if best_idx[j] != -1:
+                # best_idx[j] is the column index of the best site in this row
+                # Use that site's feature coordinates
+                best_col = best_idx[j]
+                feat_r[i, j] = feat_r[i, best_col]
+                feat_c[i, j] = feat_c[i, best_col]
+
+    return feat_r, feat_c
+
+if __name__ == '__main__':
+    # Example binary image (0 = background, 1 = foreground)
+    binary = np.array([[0, 1, 0],
+                       [1, 1, 1],
+                       [0, 1, 0]], dtype=np.int32)
+
+    feat_r, feat_c = edt_feature_transform_2d(binary)
+
+    # Compute Euclidean distances (squared) using the feature coordinates
+    # For each pixel, distance² = (i - feat_r[i,j])² + (j - feat_c[i,j])²
+    dist2 = (np.indices(binary.shape)[0] - feat_r) ** 2 + \
+            (np.indices(binary.shape)[1] - feat_c) ** 2
+    dist = np.sqrt(dist2)  # exact Euclidean distance
+
+    print(dist)