Hi there!
I am doing some DL image segmentation for 3D images. Since skimage&scipy's implementations are very memory and speed inefficient, I have been working on quite a few Numba optimised morphological functions and was planning to create a package of my own, until I saw your project. It's better not to let the community's efforts be scattered.
Though, from what I can see, your package for now only contains basic operations like dilate, erode, open, close, and smooth. There are lots of morphological operation candidates in skimage/scipy that are unoptimised for memory use or speed, could you consider them? E.g., reconstruction, distance transform, and even marker controlled watershed? I can provide my implementation of them for you to include, if that's helpful.
Also, currently there is no performance comparison tables or figures. Maybe it would be best to have them to show the advantage of using numba compared to direct skimage/scipy implementation, both in terms of speed and memory.
Thank you
Reconstruction:
def geodesic_reconstruction_by_erosion(mask, dynamic):
result = mask + dynamic
changed = True
i = 0
while changed:
changed = False
print(f"Iteration {i} of Geodesic Reconstruction...")
result, changed = reconstruct_forward_scan(mask, result, changed)
result, changed = reconstruct_backward_scan(mask, result, changed)
i += 1
return result
@njit
def reconstruct_forward_scan(mask, result, changed):
for z in range(mask.shape[0]):
for y in range(mask.shape[1]):
for x in range(mask.shape[2]):
current_value = result[z, y, x]
min_value = current_value
for dz in (-1, 0, 1):
nz = z + dz
if nz < 0 or nz >= mask.shape[0]:
continue
for dy in (-1, 0, 1):
ny = y + dy
if ny < 0 or ny >= mask.shape[1]:
continue
for dx in (-1, 0, 1):
if dz == 0 and dy == 0 and dx == 0:
continue
if dz < 0 or (dz == 0 and dy < 0) or (dz == 0 and dy == 0 and dx < 0):
nx = x + dx
if nx < 0 or nx >= mask.shape[2]:
continue
min_value = min(min_value, result[nz, ny, nx])
min_value = max(min_value, mask[z, y, x])
if min_value < current_value:
result[z, y, x] = min_value
changed = True
return result, changed
@njit
def reconstruct_backward_scan(mask, result, changed):
for z in range(mask.shape[0] - 1, -1, -1):
for y in range(mask.shape[1] - 1, -1, -1):
for x in range(mask.shape[2] - 1, -1, -1):
current_value = result[z, y, x]
min_value = current_value
for dz in (-1, 0, 1):
nz = z + dz
if nz < 0 or nz >= mask.shape[0]:
continue
for dy in (-1, 0, 1):
ny = y + dy
if ny < 0 or ny >= mask.shape[1]:
continue
for dx in (-1, 0, 1):
if dz == 0 and dy == 0 and dx == 0:
continue
if dz > 0 or (dz == 0 and dy > 0) or (dz == 0 and dy == 0 and dx > 0):
nx = x + dx
if nx < 0 or nx >= mask.shape[2]:
continue
min_value = min(min_value, result[nz, ny, nx])
min_value = max(min_value, mask[z, y, x])
if min_value < current_value:
result[z, y, x] = min_value
changed = True
return result, changed
Distance transform:
@njit(parallel=True, nogil=True)
def chamfer_distance_transform_parallel(img, dtype, num_core=16):
"""
Parallel chamfer distance transform with early stopping.
Taken from the paper "PARALLEL IMPLEMENTATION OF GEODESIC DISTANCE TRANSFORM WITH APPLICATION IN SUPERPIXEL SEGMENTATION"
"""
Z, Y, X = img.shape
num_bands = num_core
chunk_size = (Z // num_bands) + 1
max_val = np.iinfo(dtype).max
result = np.where(img > 0, dtype(max_val), dtype(0))
print(f"Need maximum {Z} Iteration for Chamfer Distance Transform")
# 3D chamfer weights: face=3, edge=4, corner=5
causal = [
(-1, -1, -1, 5), (-1, -1, 0, 4), (-1, -1, 1, 5),
(-1, 0, -1, 4), (-1, 0, 0, 3), (-1, 0, 1, 4),
(-1, 1, -1, 5), (-1, 1, 0, 4), (-1, 1, 1, 5),
(0, -1, -1, 4), (0, -1, 0, 3), (0, -1, 1, 4),
(0, 0, -1, 3),
]
anti_causal = [
(1, -1, -1, 5), (1, -1, 0, 4), (1, -1, 1, 5),
(1, 0, -1, 4), (1, 0, 0, 3), (1, 0, 1, 4),
(1, 1, -1, 5), (1, 1, 0, 4), (1, 1, 1, 5),
(0, 1, -1, 4), (0, 1, 0, 3), (0, 1, 1, 4),
(0, 0, 1, 3),
]
for it in range(Z):
print(f"Iteration {it} for Chamfer Distance Transform")
changed = np.zeros(num_bands, dtype=np.bool_) # reset for this iteration
# ---- Forward pass (parallel over bands) ----
for band in prange(num_bands):
start_z = band * chunk_size
end_z = min(start_z + chunk_size, Z)
for z in range(start_z, end_z):
for y in range(Y):
for x in range(X):
if img[z, y, x] == 0:
continue
cur = result[z, y, x]
new_val = cur
for dz, dy, dx, w in causal:
nz = z + dz
ny = y + dy
nx = x + dx
if (0 <= nz < Z and 0 <= ny < Y and 0 <= nx < X):
neigh = result[nz, ny, nx]
if neigh != max_val:
cand = neigh + w
if cand < new_val:
new_val = cand
if new_val < cur:
result[z, y, x] = new_val
changed[band] = True # record that this band was updated
# ---- Backward pass (parallel over bands) ----
for band in prange(num_bands):
start_z = band * chunk_size
end_z = min(start_z + chunk_size, Z)
for z in range(end_z - 1, start_z - 1, -1):
for y in range(Y - 1, -1, -1):
for x in range(X - 1, -1, -1):
if img[z, y, x] == 0:
continue
cur = result[z, y, x]
new_val = cur
for dz, dy, dx, w in anti_causal:
nz = z + dz
ny = y + dy
nx = x + dx
if (0 <= nz < Z and 0 <= ny < Y and 0 <= nx < X):
neigh = result[nz, ny, nx]
if neigh != max_val:
cand = neigh + w
if cand < new_val:
new_val = cand
if new_val < cur:
result[z, y, x] = new_val
changed[band] = True # record that this band was updated
# After both passes, check if any change occurred anywhere
if not np.any(changed):
print("No more change detected! Stopping Early.")
break # converged
return result
Marker controlled watershed:
def __heapify_markers_3d(markers, image):
"""Create a priority queue heap with the markers on it for 3D."""
stride = np.array(image.strides, dtype=np.uint32) // image.itemsize
coords = np.argwhere(markers != 0).astype(np.uint32)
ncoords = coords.shape[0]
if ncoords > 0:
pixels = image[markers != 0]
age = np.arange(ncoords, dtype=np.uint32)
offset = np.zeros(coords.shape[0], dtype=np.uint32)
for i in range(image.ndim):
offset = offset + stride[i] * coords[:, i]
pq = [tuple(row) for row in np.column_stack((pixels, age, offset, coords))]
ordering = np.lexsort((age, pixels))
pq = [pq[i] for i in ordering]
else:
pq = np.zeros((0, markers.ndim + 3), int)
return (pq, ncoords)
@njit(nogil=True)
def _watershed_loop(pq, labels, connect_increments, mask, image, age):
max_x, max_y, max_z = labels.shape
total_pixels = image.size if mask is None else np.count_nonzero(mask)
processed = 0
print_interval = max(1, total_pixels // 20) # print every 5%
print(f"A total of {total_pixels} needs to be processed.")
while len(pq):
pix_value, pix_age, _, pix_x, pix_y, pix_z = heappop(pq)
processed += 1
pix_label = labels[pix_x, pix_y, pix_z]
if processed % print_interval == 0:
progress = processed * 100 // total_pixels
print(f"Watershed Progress: {progress}% ({processed}/{total_pixels})")
for dx, dy, dz in connect_increments:
x, y, z = pix_x + dx, pix_y + dy, pix_z + dz
if x < 0 or y < 0 or z < 0 or x >= max_x or y >= max_y or z >= max_z:
continue
if labels[x, y, z]:
continue
if mask is not None and not mask[x, y, z]:
continue
labels[x, y, z] = pix_label
new_pq_item = (np.uint32(image[x, y, z]), np.uint32(age), np.uint32(0), np.uint32(x), np.uint32(y), np.uint32(z))
heappush(pq, new_pq_item)
age += 1
return labels
def marker_controlled_watershed(image, markers, mask=None):
connect_increments = [
(1, 0, 0), (-1, 0, 0), (0, 1, 0), (0, -1, 0), (0, 0, 1), (0, 0, -1)
]
print("Starts watershed flooding...")
pq, age = __heapify_markers_3d(markers, image)
return _watershed_loop(pq, markers, connect_increments, mask, image, age)
Hi there!
I am doing some DL image segmentation for 3D images. Since skimage&scipy's implementations are very memory and speed inefficient, I have been working on quite a few Numba optimised morphological functions and was planning to create a package of my own, until I saw your project. It's better not to let the community's efforts be scattered.
Though, from what I can see, your package for now only contains basic operations like dilate, erode, open, close, and smooth. There are lots of morphological operation candidates in skimage/scipy that are unoptimised for memory use or speed, could you consider them? E.g., reconstruction, distance transform, and even marker controlled watershed? I can provide my implementation of them for you to include, if that's helpful.
Also, currently there is no performance comparison tables or figures. Maybe it would be best to have them to show the advantage of using numba compared to direct skimage/scipy implementation, both in terms of speed and memory.
Thank you
Reconstruction:
Distance transform:
Marker controlled watershed: