diff --git a/dither.py b/dither.py index 23d6f79..ebd2637 100644 --- a/dither.py +++ b/dither.py @@ -55,11 +55,16 @@ def rgb_to_linear(rgb): PICO_RGB_LINEAR = [rgb_to_linear(x) for x in PICO_RGB] -bayer = np.array([[0]]) -np.block([[4*bayer, 4*bayer+3],[4*bayer+2, 4*bayer+1]]) -bayer = np.block([[4*bayer, 4*bayer+3],[4*bayer+2, 4*bayer+1]]) -bayer = np.block([[4*bayer, 4*bayer+3],[4*bayer+2, 4*bayer+1]]) -bayer = np.block([[4*bayer, 4*bayer+3],[4*bayer+2, 4*bayer+1]]) +def make_bayer_matrix(size): + assert size > 0 + assert size & (size - 1) == 0 # power of two + m = np.array([[0.0]]) + while m.shape[0] < size: + m = np.block([[4*m, 4*m+3],[4*m+2, 4*m+1]]) + return m + +bayer = make_bayer_matrix(4) + def dither_pixel(img, xy, pal): global bayer pal2 = [np.array(rgb_to_linear(c)) for c in pal] @@ -72,12 +77,12 @@ def dither_pixel(img, xy, pal): return index color = np.array(rgb_to_linear(img.getpixel(xy))) err = np.array([0.0,0.0,0.0]) - for _ in range(64): + for _ in range(bayer.size): err += color i = nearest(err) counts[i] += 1 err -= pal2[i] - thresh = bayer[y % 8][x % 8] + thresh = bayer[y % bayer.shape[0]][x % bayer.shape[1]] for i, count in enumerate(counts): thresh -= count if thresh < 0: