map_labeling.py

from type import array_t

import numpy as np_


invalid_n_neighbors_c = 27  # Must be positive and higher (strictly) than
# the max number of neighbors in a skeleton

shifts_of_26_neighbors_c = tuple(
    (i, j, k)
    for i in (-1, 0, 1)
    for j in (-1, 0, 1)
    for k in (-1, 0, 1)
    if i != 0 or j != 0 or k != 0
)


def PartLMap(map_: array_t) -> array_t:
    #
    # The part mask is labeled as follows:
    # background=invalid_n_neighbors_c; Pixels of the skeleton=number of
    # neighboring pixels that belong to the skeleton (as expected,
    # isolated pixels receive 0).
    #
    result = np_.array(map_, dtype=np_.int8)
    result[result > 0] = 1
    padded_sm = np_.pad(map_ > 0, 1, "constant")

    for shifts in shifts_of_26_neighbors_c:
        result += np_.roll(padded_sm, shifts, axis=(0, 1, 2))[1:-1, 1:-1, 1:-1]

    result[map_ == 0] = invalid_n_neighbors_c + 1

    return result - 1