Ordered Dithering

ordered_dithering.py

@functools.lru_cache()
def dither_matrix(n: int):
    assert isinstance(n, int) and (n > 0)
    if n == 1:
        return np.array([[0]])
    matrix = (1/n**2) * np.asarray(np.bmat([
        [n**2 * dither_matrix(int(n/2)) + 0, n**2 * dither_matrix(int(n/2)) + 2],
        [n**2 * dither_matrix(int(n/2)) + 3, n**2 * dither_matrix(int(n/2)) + 1],
    ]))
    matrix.flags.writeable = False  # because we are caching this
    return matrix

def ordered_dithering(img: np.ndarray, n: int):
    d_mat = dither_matrix(n=n)
    # get indices
    yy = np.arange(img.shape[0])[:, None]  # (H, 1)
    xx = np.arange(img.shape[1])[None, :]  # (1, W)
    # get values
    dd = d_mat[yy % d_mat.shape[0], xx % d_mat.shape[0]]  # (H, W)
    # add dims
    if img.ndim == 3:
        dd = dd[..., None]
    # dither
    return (img > dd)