python taichi 加速 dither仿色抖动算法

教程

9种dither算法与历史发展
最开始是如何生成bayer矩阵
wiki: bayer有序抖动
python生成任意规模bayer矩阵
知乎：dither启发的艺术效果，半调/柱形
taichi ndarray文档
CSDN dither算法

理解bayer矩阵

因为bayer矩阵要平铺，所以要分治、递归地解决。同时，分治也能保证像素之间平等的稀疏。而分治，就限制了n为2的幂次方。
我想生成任意大小的bayer矩阵，比如5x5、7x7、11x11，但网上似乎都没有教程。

胡言乱语2

在胡言乱语1后，又研究了下，发现bayer矩阵跟自己随意捏的矩阵相比，有其特殊性：

• 稀疏性
• 一些阶段能看出是一条斜线

就像魔方与多元可微一样，不同维度之间有耦合。

胡言乱语1

按照这个博客的说法：最开始是如何生成bayer矩阵，会有点难以理解，为什么bayer矩阵要按这个顺序

A 2×2 Bayer matrix looks like this, and is constructed by starting anywhere, going to the furthest pixel from the starting point (which is a diagonal line), and then filling the remaining two pixels by the same logic.
2×2 拜耳矩阵如下所示，其构造方式为：从任意位置开始，到距起点最远的像素（即对角线），然后按相同的逻辑填充其余两个像素。

其实，如果将bayer矩阵平铺一下，我们能得到类似行列式的规律：

这样，3x3的bayer矩阵也能解释得通了。平铺之后，都是以↗的顺序来蒙版遮罩。

当n较大时，甚至可以自定义bayer矩阵的顺序，只需保证里面的数字不重复，从0递增到n*n。

代码实现

taichi_dither.py

#!/bin/env python
import taichi as ti
import numpy as np
import cv2
from copy import deepcopy
ti.init(arch=ti.cpu)
 
DEBUG=True
MAX=255
np.set_printoptions(threshold=np.inf, linewidth=180)  # numpy打印选项
img_from='/home/n/photo/Portal_Companion_Cube.jpg'
# img_from='/home/n/photo/neco godness.jpg'
# img_from='/media/n/data/download/firefox/updated/browser/chrome/icons/default/default32.png'
 
def show_image(img):
    if isinstance(img, str):
        img = cv2.imread(img)
    elif img.dtype != 'uint8':
        img = np.clip(img, 0, MAX)  # 将图片像素值限制在 0~255 之间
        img = img.astype(np.uint8)
    
    # 显示图片
    cv2.imshow('Image', img)
    while True:
        key = cv2.waitKey(1) & 0xFF
        if key == 27:  # 27 是 ESC 键的 ASCII 码
            break
    cv2.destroyAllWindows()
 
 
PALETTE=[
0x00,0xFF
]
# img2d=ti.types.ndarray(dtype=ti.math.vec3, ndim=2)
type_img2d = ti.types.ndarray(element_dim=0,ndim=2)
type_bayerM = ti.types.ndarray(element_dim=0,ndim=2)
 
@ti.func
def clamp(x,min=-MAX,max=MAX):
    """控制出血阈值，建议min in [-256,0]"""
    return ti.math.clamp(x,min,max)
    # return np.clip(x,min,max)
 
@ti.kernel
def dither_basic(img:type_img2d):
    h,w = img.shape
    for i in range(h):
      for j in range(w):
        min_distance = MAX
        oldpixel = img[i, j]
        newpixel = 0
        for c in ti.static(PALETTE):
            distance = abs(img[i, j] - c)
            if distance < min_distance:
                min_distance = distance
                newpixel = c
        img[i, j] = newpixel
        if j + 1 < w:
            img[i, j+1] += oldpixel - newpixel
 
@ti.kernel
def dither_floyd(img:type_img2d):
    h,w = img.shape
    for i in range(h):
      for j in range(w):
    # for i,j in ti.ndrange(h,w):  # taichi的ti.ndrange有bug，与下面的结果不同！
        oldpixel = img[i, j]
        newpixel = 0 if oldpixel < 128 else MAX
        img[i, j] = newpixel
        quant_error = oldpixel - newpixel
        if j + 1 < img.shape[1]:
            img[i, j + 1] += quant_error * 7 >> 4
        if i + 1 < img.shape[0]:
            if j - 1 >= 0:
                img[i + 1, j - 1] += quant_error * 3 >> 4
            img[i + 1, j] += quant_error * 5 >> 4
 
def bit_reverse(x, n):
    return int(bin(x)[2:].zfill(n)[::-1], 2)
 
def bit_interleave(x, y, n):
    x = bin(x)[2:].zfill(n)
    y = bin(y)[2:].zfill(n)
    return int(''.join(''.join(i) for i in zip(x, y)), 2)
 
def bayer_entry(x, y, n):
    return bit_reverse(bit_interleave(x ^ y, y, n), 2*n)
 
def bayer_matrix(n):
    """https://gamedev.stackexchange.com/questions/130696/how-to-generate-bayer-matrix-of-arbitrary-size"""
    r = range(2**n)
    return [[bayer_entry(x, y, n) for x in r] for y in r]
 
@ti.kernel
def dither_bayer(img:type_img2d, bayerM:type_bayerM):
    h,w = img.shape
    n = bayerM.shape[0]
    for i in range(h):
      for j in range(w):
        threshold = bayerM[i % n, j % n] * MAX // (n**2)
        if img[i, j] > threshold:
            img[i, j] = MAX
        else:
            img[i, j] = 0
 
@ti.kernel
def dither_atkinson(img:type_img2d):
    h,w = img.shape
    for i in range(h):
      for j in range(w):
        oldpixel = img[i, j]
        newpixel = 0 if oldpixel < 128 else MAX
        img[i, j] = newpixel
        quant_err = oldpixel - newpixel
 
        if j + 1 < img.shape[1]:
            img[i, j + 1] += quant_err >> 3
        if j + 2 < img.shape[1]:
            img[i, j + 2] += quant_err >> 3
        if i + 1 < img.shape[0]:
            if j - 1 >= 0:
                img[i + 1, j - 1] += quant_err >> 3
            img[i + 1, j] += quant_err >> 3
            if j + 1 < img.shape[1]:
                img[i + 1, j + 1] += quant_err >> 3
            if j + 2 < img.shape[1]:
                img[i + 1, j + 2] += quant_err >> 3
        if i + 2 < img.shape[0]:
            if j - 1 >= 0:
                img[i + 2, j - 1] += quant_err >> 3
            img[i + 2, j] += quant_err >> 3
            if j + 1 < img.shape[1]:
                img[i + 2, j + 1] += quant_err >> 3
            if j + 2 < img.shape[1]:
                img[i + 2, j + 2] += quant_err >> 3
 
def diff(img1,img2):
    return np.sum(np.abs(img1-img2))
 
img = cv2.imread(img_from)
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
img_u8 = deepcopy(img) if DEBUG else None
img = img.astype(np.int32)  # 转换为 int32 类型
 
if DEBUG:
    cmd = 'diff(img,img_u8)'
    print('i32<-u8',cmd,'=', eval(cmd))
 
    # print('Raw:')
    # print(img, '\n')
 
n=4
bayerM = ti.ndarray(shape=(2**n, 2**n), dtype=ti.u8)
bayerM.from_numpy(np.matrix(bayer_matrix(n), dtype=np.uint8))
dither_bayer(img, bayerM)
if DEBUG:
    img_i32 = deepcopy(img)
 
    # print('dithered:',diff(img,img_u8),np.max(img), np.min(img), img.shape)
    # print(img, '\n')
 
if DEBUG:
    img = np.clip(img, 0, MAX)
    img = img.astype(np.uint8)
    cmd = 'diff(img,img_i32)'
    print('u8<-i32',cmd,'=', eval(cmd),np.max(img), np.min(img))
 
    # print(img, '\n')
show_image(img)
 
 
 
# neighbors = [(-1, -1), (-1, ), (-1, +1),
#              (0, -1),             (0, +1),
#              (+1, -1), (+1, ), (+1, +1)] # 简化写法
 
# values = [1, 2, 3, 4, 5, 6, 7, 8]  # 示例值
# # a = np.arange(100).reshape(10, 10)
# a = np.zeros((10, 10))
# print(a,'\n')
# for i in range(10-3):
#   for j in range(10-3):
#     for (x, y), value in zip(neighbors, values):
#         a[i+x, j+y] = value
# print(a,'\n')