The is so useful, thank you for making this.
May I suggest one improvement: adding a small zero padding / zero stripping utility.
Because currently, if a connected component is at the edge, looking for it's out of bound neighbors will result an out of bound index error. Instead of checking for index, maybe simply padding the 3d volume with a 'Null' value will be a more versatile solution.
def zero_pad_3d(np_3d):
return np.pad(np_3d, ((1,1), (1,1), (1,1)), 'constant')
def zero_strip_3d(np_3d):
return np_3d[1:-1, 1:-1, 1:-1]
Then after calculations, strip away the paddings to return back the original size.