labeling function using simple flood fill algorithm

labels.py

import numpy as np
from numba import jit
from typing import Sequence

# @jit(nopython=True)
def flood_fill_labels(data: np.ndarray, seeds: Sequence[Sequence[int]]=[]) -> np.ndarray:
    """Labels an array using a flood fill algorithm

    The function implements a flooding algorithm 
    considering a connectivity equals to the 2*data.ndim

    Parameters
    ----------
        data : array_like
            Multidimensional array conatining the data.
        seeds : array_like
            n-by-d array with `n` seeds,
            `d` must be equal to `data.ndim`.

    Return
    ------
        labels : array (numpy)
            label array with the same shape as `data`.

    """
    # parse seeds
    seeds_ = set() # use set remove duplicates
    if not seeds:
        for seed in list(zip(*np.where(data))):
            seeds_.add(tuple(seed))
    else:
        for seed in list(seeds):
            seeds_.add(tuple(seed))
    # create an empty array from data shape
    labels = np.zeros_like(data, dtype=np.int32)
    # store array shape (used later)
    ndims = data.shape
    # create a labels counter
    label_count = 0
    # for each seed
    for seed in seeds_:
        # skip seeds in a previously labeled blob
        if labels[seed]:
            continue
        # update label_count
        label_count += 1
        # get the reference value
        ref = data[seed]
        # create a todo list
        todo = [seed]
        # while the todo list is not empty
        while todo:
            # pop the last todo input : pixel to process
            current = todo.pop()
            # if already flooded pass to the next, else flood (if condition is satisfied)
            if not labels[current] and data[current] == ref:
                labels[current] = label_count
                # after a flood, add neighbors to the todo list (propagate the flooding)
                for n, d in enumerate(current):  # works with ND arrays
                    if d > 0:  # check if lower border
                        indexes = list(current)
                        indexes[n] -= 1
                        todo.append(tuple(indexes))
                    if d < ndims[n] - 1:  # check if upper border
                        indexes = list(current)
                        indexes[n] += 1
                        todo.append(tuple(indexes))
    # the end
    return labels


if __name__ == "__main__":
    data = np.asarray([
        [0, 0, 0, 0, 0, 1, 0],
        [0, 1, 0, 0, 1, 0, 0],
        [0, 0, 0, 0, 1, 1, 0],
        [0, 1, 1, 0, 0, 0, 0]], dtype=int)
    seeds = [(1, 1), (1, 4), (3, 1)]
    labels = flood_fill_labels(data, seeds)
    print('data:\n{}\nseeds:\n{}\nlabels:\n{}'.format(
        data, seeds, labels))