def should_highlight_pixel(arr2d, set1d):
    """Compute boolean mask for items in arr2d that are also in set1d"""
    H, W = arr2d.shape
    output = np.zeros((H, W), dtype=bool_)
    for i in range(H):
        for j in range(W):
            for n in range(set1d.size):
                output[i, j] = output[i, j] or (arr2d[i, j] == set1d[n])
    return output
Compute boolean mask for items in arr2d that are also in set1d
infinigen/tools/ground_truth/segmentation_lookup.py