Resolve Mask Occlusions #python

resolve_mask_occlusions.py

from sortedcontainers import SortedDict

def resolve_mask_occlusions(mask, z_index):
    if (len(mask.shape) != 3 or mask.shape[2] != len(z_index)):
        raise TypeError('Mask shape depth must match the z_index')
    z_index_map = SortedDict()
    for mask_index, mask_depth in enumerate(z_index):
        if mask_depth in z_index_map:
            z_index_map[mask_depth].append(mask_index)
        else:
            z_index_map[mask_depth] = [mask_index]
    # list of mask layer indexes ordered from foreground to background
    depth_list = reversed(z_index_map.values())
    # initial combined mask and combined occlusion
    layer_indexes = next(depth_list)
    combined_mask = mask[:, :, layer_indexes.pop()].copy()
    for index in layer_indexes:
        combined_mask |= mask[:, :, index]
    combined_occlusion = combined_mask ^ 1
    # iterate to the background layer
    for layer_indexes in depth_list:
        index = layer_indexes.pop()
        mask[:, :, index] *= combined_occlusion
        combined_mask = mask[:, :, index].copy()
        for index in layer_indexes:
            mask[:, :, index] *= combined_occlusion
            combined_mask |= mask[:, :, index]
        combined_occlusion &= combined_mask ^ 1

# mask is a numpy arrwy with shape: [Height, Width, Z-Index]
# z_index should be a list of relative depths e.g. [0,1,1,-1,2,5] ordered by the Z-Index of the mask
# high numbers means near the foreground, low numbers means near the background
# the highest number is the foreground
# the lowest number is the background
# it possible for 2 or more layers of the mask to have the same depth
# this means they will not occlude each other, but their combined mask will occlude lower layers