Hi Everyone:I had implemented a watershed by numba, and now it has pass all the
test_watershed check and got a exactly same result with skimage, support
watershed-line, support connectivity, (pass test 0-12, but mine did not support
compact yet)
It support nd, now only support uint8 image with uint16 marks. (because I think
sometimes it is enough, It can process float image if it is needed)
And it is very smart and fast. less than 100 lines, and faster than skimage
one. a 3d image which skimage run more than 10 minutes ( and I cannot wait any
longer, so kill it), but my watershed just need 28 s.
The code are attached.
BestYXDragon
import numpy as np
from numba import jit
from scipy.ndimage import label, generate_binary_structure
from scipy.misc import imread, imsave
def neighbors(shape, conn=1):
dim = len(shape)
block = generate_binary_structure(dim, conn)
block[tuple([1]*dim)] = 0
idx = np.where(block>0)
idx = np.array(idx, dtype=np.uint8).T
idx = np.array(idx-[1]*dim)
acc = np.cumprod((1,)+shape[::-1][:-1])
return np.dot(idx, acc[::-1])
@jit
def step(img, msk, line, pts, s, level, up, nbs):
cur = 0
while cur<s:
p = pts[cur]
if msk[p] == 0xfffe or up and img[p]>level or not up and img[p]<level:
cur += 1
continue
for dp in nbs:
cp = p+dp
if msk[cp]==0:
msk[cp] = msk[p]
if s == len(pts):
s, cur = clear(pts, s, cur)
pts[s] = cp
s+=1
elif msk[cp] == msk[p]:
continue
elif msk[cp] == 0xffff:
continue
elif msk[cp] == 0xfffe:
continue
elif line : msk[cp] = 0xfffe
pts[cur] = -1
cur+=1
return cur
@jit
def clear(pts, s, cur):
ns = 0; nc=0;
for c in range(s):
if pts[c]!=-1:
pts[ns] = pts[c]
ns += 1
if c<cur:nc += 1
return ns, nc
@jit
def collect(img, mark, nbs, pts):
bins = np.zeros(img.max()+1, dtype=np.uint32)
cur = 0
for p in range(len(mark)):
bins[img[p]] += 1
if mark[p]==0xffff: continue # edge
if mark[p]==0: continue # zero
for dp in nbs:
if mark[p+dp]!=mark[p]:
pts[cur] = p
cur += 1
break
return cur, bins
@jit
def erose(mark):
for i in range(len(mark)):
if mark[i]>0xfff0:mark[i]=0
def buffer(img, dtype):
buf = np.zeros(tuple(np.array(img.shape)+2), dtype=dtype)
buf[tuple([slice(1,-1)]*buf.ndim)] = img
return buf
def watershed(img, mark, conn=1, line=False, up=True):
maxv = img.max(); minv = img.min();
if img.dtype != np.uint8:
img = ((img-minv)*(255/(maxv-minv))).astype(np.uint8)
img = buffer(img, np.uint8)
mark = buffer(mark, np.uint16)
ndim = img.ndim
for n in range(ndim):
idx = [slice(None) if i!=n else [0,-1] for i in range(ndim)]
mark[tuple(idx)] = 0xffff
nbs = neighbors(img.shape, conn)
img = img.ravel()
mark1d = mark.ravel()
pts = np.zeros(img.size//3, dtype=np.int64)
s, bins = collect(img, mark1d, nbs, pts)
for level in range(len(bins))[::1 if up else -1]:
if bins[level]==0:continue
s, c = clear(pts, s, 0)
s = step(img, mark1d, line, pts, s, level, up, nbs)
if line:erose(mark1d)
return mark[tuple([slice(1,-1)]*ndim)]
#============================= test =============================
from skimage.morphology.watershed import watershed as skwatershed
def test_watershed01():
"watershed 1"
data = np.array([[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 1, 1, 1, 1, 1, 0],
[0, 1, 0, 0, 0, 1, 0],
[0, 1, 0, 0, 0, 1, 0],
[0, 1, 0, 0, 0, 1, 0],
[0, 1, 1, 1, 1, 1, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0]], np.uint8)
markers = np.array([[ 2, 0, 0, 0, 0, 0, 0],
[ 0, 0, 0, 0, 0, 0, 0],
[ 0, 0, 0, 0, 0, 0, 0],
[ 0, 0, 0, 0, 0, 0, 0],
[ 0, 0, 0, 0, 0, 0, 0],
[ 0, 0, 0, 1, 0, 0, 0],
[ 0, 0, 0, 0, 0, 0, 0],
[ 0, 0, 0, 0, 0, 0, 0],
[ 0, 0, 0, 0, 0, 0, 0],
[ 0, 0, 0, 0, 0, 0, 0]], np.uint16)
out1 = watershed(data, markers)
out2 = skwatershed(data, markers)
print(np.sum(out1-out2!=0))
def test_watershed02():
"watershed 2"
data = np.array([[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 1, 1, 1, 1, 1, 0],
[0, 1, 0, 0, 0, 1, 0],
[0, 1, 0, 0, 0, 1, 0],
[0, 1, 0, 0, 0, 1, 0],
[0, 1, 1, 1, 1, 1, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0]], np.uint8)
markers = np.array([[2, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 1, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0]], np.uint16)
out1 = watershed(data, markers, 2)
out2 = skwatershed(data, markers, 2)
print(np.sum(out1-out2!=0))
def test_watershed03():
"watershed 3"
data = np.array([[0, 0, 0, 0, 0, 0, 0],
[0, 1, 1, 1, 1, 1, 0],
[0, 1, 0, 1, 0, 1, 0],
[0, 1, 0, 1, 0, 1, 0],
[0, 1, 0, 1, 0, 1, 0],
[0, 1, 1, 1, 1, 1, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0]], np.uint8)
markers = np.array([[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 2, 0, 3, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 1]], np.uint16)
out1 = watershed(data, markers)
out2 = skwatershed(data, markers)
print(np.sum(out1-out2!=0))
def test_watershed04():
"watershed 4"
data = np.array([[0, 0, 0, 0, 0, 0, 0],
[0, 1, 1, 1, 1, 1, 0],
[0, 1, 0, 1, 0, 1, 0],
[0, 1, 0, 1, 0, 1, 0],
[0, 1, 0, 1, 0, 1, 0],
[0, 1, 1, 1, 1, 1, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0]], np.uint8)
markers = np.array([[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 2, 0, 3, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 1]], np.uint8)
out1 = watershed(data, markers, 2)
out2 = skwatershed(data, markers, 2)
print(np.sum(out1-out2!=0))
def test_watershed05():
"watershed 4"
data = np.array([[0, 0, 0, 0, 0, 0, 0],
[0, 1, 1, 1, 1, 1, 0],
[0, 1, 0, 1, 0, 1, 0],
[0, 1, 0, 1, 0, 1, 0],
[0, 1, 0, 1, 0, 1, 0],
[0, 1, 1, 1, 1, 1, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0]], np.uint8)
markers = np.array([[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 3, 0, 2, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 1]], np.uint8)
out1 = watershed(data, markers, 2)
out2 = skwatershed(data, markers, 2)
print(np.sum(out1-out2!=0))
def test_watershed06():
"watershed 6"
data = np.array([[0, 1, 0, 0, 0, 1, 0],
[0, 1, 0, 0, 0, 1, 0],
[0, 1, 0, 0, 0, 1, 0],
[0, 1, 1, 1, 1, 1, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0]], np.uint8)
markers = np.array([[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 1, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[2, 0, 0, 0, 0, 0, 0]], np.uint16)
out1 = watershed(data, markers, 2)
out2 = skwatershed(data, markers, 2)
print(np.sum(out1-out2!=0))
def test_watershed07():
"A regression test of a competitive case that failed"
data = np.array([[255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
255, 255, 255, 255, 255],
[255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
255, 255, 255, 255, 255],
[255, 255, 255, 255, 255, 204, 204, 204, 204, 204, 204,
255, 255, 255, 255, 255],
[255, 255, 255, 204, 204, 183, 153, 153, 153, 153, 183,
204, 204, 255, 255, 255],
[255, 255, 204, 183, 153, 141, 111, 103, 103, 111, 141,
153, 183, 204, 255, 255],
[255, 255, 204, 153, 111, 94, 72, 52, 52, 72, 94,
111, 153, 204, 255, 255],
[255, 255, 204, 153, 111, 72, 39, 1, 1, 39, 72,
111, 153, 204, 255, 255],
[255, 255, 204, 183, 141, 111, 72, 39, 39, 72, 111,
141, 183, 204, 255, 255],
[255, 255, 255, 204, 183, 141, 111, 72, 72, 111, 141,
183, 204, 255, 255, 255],
[255, 255, 255, 255, 204, 183, 141, 94, 94, 141, 183,
204, 255, 255, 255, 255],
[255, 255, 255, 255, 255, 204, 153, 103, 103, 153, 204,
255, 255, 255, 255, 255],
[255, 255, 255, 255, 204, 183, 141, 94, 94, 141, 183,
204, 255, 255, 255, 255],
[255, 255, 255, 204, 183, 141, 111, 72, 72, 111, 141,
183, 204, 255, 255, 255],
[255, 255, 204, 183, 141, 111, 72, 39, 39, 72, 111,
141, 183, 204, 255, 255],
[255, 255, 204, 153, 111, 72, 39, 1, 1, 39, 72,
111, 153, 204, 255, 255],
[255, 255, 204, 153, 111, 94, 72, 52, 52, 72, 94,
111, 153, 204, 255, 255],
[255, 255, 204, 183, 153, 141, 111, 103, 103, 111, 141,
153, 183, 204, 255, 255],
[255, 255, 255, 204, 204, 183, 153, 153, 153, 153, 183,
204, 204, 255, 255, 255],
[255, 255, 255, 255, 255, 204, 204, 204, 204, 204, 204,
255, 255, 255, 255, 255],
[255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
255, 255, 255, 255, 255],
[255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
255, 255, 255, 255, 255]], dtype=np.uint8)
mask = (data != 255)
markers = np.zeros(data.shape, np.uint16)
markers[6, 7] = 1
markers[14, 7] = 2
out1 = watershed(data, markers, 2)
out2 = skwatershed(data, markers, 2)
print(np.sum(out1-out2!=0))
def test_watershed08():
"The border pixels + an edge are all the same value"
data = np.array([[255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
255, 255, 255, 255, 255],
[255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
255, 255, 255, 255, 255],
[255, 255, 255, 255, 255, 204, 204, 204, 204, 204, 204,
255, 255, 255, 255, 255],
[255, 255, 255, 204, 204, 183, 153, 153, 153, 153, 183,
204, 204, 255, 255, 255],
[255, 255, 204, 183, 153, 141, 111, 103, 103, 111, 141,
153, 183, 204, 255, 255],
[255, 255, 204, 153, 111, 94, 72, 52, 52, 72, 94,
111, 153, 204, 255, 255],
[255, 255, 204, 153, 111, 72, 39, 1, 1, 39, 72,
111, 153, 204, 255, 255],
[255, 255, 204, 183, 141, 111, 72, 39, 39, 72, 111,
141, 183, 204, 255, 255],
[255, 255, 255, 204, 183, 141, 111, 72, 72, 111, 141,
183, 204, 255, 255, 255],
[255, 255, 255, 255, 204, 183, 141, 94, 94, 141, 183,
204, 255, 255, 255, 255],
[255, 255, 255, 255, 255, 204, 153, 141, 141, 153, 204,
255, 255, 255, 255, 255],
[255, 255, 255, 255, 204, 183, 141, 94, 94, 141, 183,
204, 255, 255, 255, 255],
[255, 255, 255, 204, 183, 141, 111, 72, 72, 111, 141,
183, 204, 255, 255, 255],
[255, 255, 204, 183, 141, 111, 72, 39, 39, 72, 111,
141, 183, 204, 255, 255],
[255, 255, 204, 153, 111, 72, 39, 1, 1, 39, 72,
111, 153, 204, 255, 255],
[255, 255, 204, 153, 111, 94, 72, 52, 52, 72, 94,
111, 153, 204, 255, 255],
[255, 255, 204, 183, 153, 141, 111, 103, 103, 111, 141,
153, 183, 204, 255, 255],
[255, 255, 255, 204, 204, 183, 153, 153, 153, 153, 183,
204, 204, 255, 255, 255],
[255, 255, 255, 255, 255, 204, 204, 204, 204, 204, 204,
255, 255, 255, 255, 255],
[255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
255, 255, 255, 255, 255],
[255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
255, 255, 255, 255, 255]])
mask = (data != 255)
markers = np.zeros(data.shape, int)
markers[6, 7] = 1
markers[14, 7] = 2
out1 = watershed(data, markers, 2)
out2 = skwatershed(data, markers, 2)
print(np.sum(out1-out2!=0))
def test_watershed09():
'''this is a float test'''
pass
def test_watershed10():
"watershed 10"
data = np.array([[1, 1, 1, 1],
[1, 1, 1, 1],
[1, 1, 1, 1],
[1, 1, 1, 1]], np.uint8)
markers = np.array([[1, 0, 0, 2],
[0, 0, 0, 0],
[0, 0, 0, 0],
[3, 0, 0, 4]], np.int8)
out1 = watershed(data, markers, 2)
out2 = skwatershed(data, markers, 2)
print(np.sum(out1-out2!=0))
def test_watershed11():
'''Make sure that all points on this plateau are assigned to closest seed'''
data = np.zeros((21, 21), dtype=np.uint8)
markers = np.zeros((21, 21), int)
markers[5, 5] = 1
markers[5, 10] = 2
markers[10, 5] = 3
markers[10, 10] = 4
structure = np.array([[False, True, False],
[True, True, True],
[False, True, False]])
out1 = watershed(data, markers)
out2 = skwatershed(data, markers, structure)
print(np.sum(out1-out2!=0))
def test_watershed12():
"The watershed line"
data = np.array([[203, 255, 203, 153, 153, 153, 153, 153, 153, 153, 153,
153, 153, 153, 153, 153],
[203, 255, 203, 153, 153, 153, 102, 102, 102, 102, 102,
102, 153, 153, 153, 153],
[203, 255, 203, 203, 153, 153, 102, 102, 77, 0, 102,
102, 153, 153, 203, 203],
[203, 255, 255, 203, 153, 153, 153, 102, 102, 102, 102,
153, 153, 203, 203, 255],
[203, 203, 255, 203, 203, 203, 153, 153, 153, 153, 153,
153, 203, 203, 255, 255],
[153, 203, 255, 255, 255, 203, 203, 203, 203, 203, 203,
203, 203, 255, 255, 203],
[153, 203, 203, 203, 255, 255, 255, 255, 255, 255, 255,
255, 255, 255, 203, 203],
[153, 153, 153, 203, 203, 203, 203, 203, 255, 203, 203,
203, 203, 203, 203, 153],
[102, 102, 153, 153, 153, 153, 203, 203, 255, 203, 203,
255, 203, 153, 153, 153],
[102, 102, 102, 102, 102, 153, 203, 255, 255, 203, 203,
203, 203, 153, 102, 153],
[102, 51, 51, 102, 102, 153, 203, 255, 203, 203, 153,
153, 153, 153, 102, 153],
[ 77, 51, 51, 102, 153, 153, 203, 255, 203, 203, 203,
153, 102, 102, 102, 153],
[ 77, 0, 51, 102, 153, 203, 203, 255, 203, 255, 203,
153, 102, 51, 102, 153],
[ 77, 0, 51, 102, 153, 203, 255, 255, 203, 203, 203,
153, 102, 0, 102, 153],
[102, 0, 51, 102, 153, 203, 255, 203, 203, 153, 153,
153, 102, 102, 102, 153],
[102, 102, 102, 102, 153, 203, 255, 203, 153, 153, 153,
153, 153, 153, 153, 153]], dtype=np.uint8)
markerbin = (data==0)
from skimage.measure import label
markers = label(markerbin).astype(np.uint16)
out1 = watershed(data, markers, 1, True)
out2 = skwatershed(data, markers, 1, watershed_line=True)
print(np.sum(out1-out2!=0))
def test_compact_watershed():
'''not support compact'''
pass
if __name__ == '__main__':
test_watershed01()
test_watershed02()
test_watershed03()
test_watershed04()
test_watershed05()
test_watershed06()
test_watershed07()
test_watershed08()
test_watershed09()
test_watershed10()
test_watershed11()
test_watershed12()
_______________________________________________
scikit-image mailing list
scikit-image@python.org
https://mail.python.org/mailman/listinfo/scikit-image
