165 lines
5.7 KiB
Python
165 lines
5.7 KiB
Python
"""图像管线：调整 → 抖动量化 → 1bpp 平面打包。
|
||
算法与参数证据：analysis/web/timedjs.js（见 docs/protocol.md §5）。
|
||
所有实现与 Web 端 JS 语义对齐：同样的灰度系数、误差核、BCD 无 —— 像素级等价为目标。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
from PIL import Image
|
||
|
||
GRAY = (0.299, 0.587, 0.114)
|
||
|
||
BAYER4 = (
|
||
(15, 135, 45, 165),
|
||
(195, 75, 225, 105),
|
||
(60, 180, 30, 150),
|
||
(240, 120, 210, 90),
|
||
)
|
||
|
||
# 误差扩散核：(dx, dy, 权重分子)；除数见 _KERNELS 键名
|
||
_KERNELS = {
|
||
"floydsteinberg": ([(1, 0, 7), (-1, 1, 3), (0, 1, 5), (1, 1, 1)], 16),
|
||
"atkinson": (
|
||
[(1, 0, 1), (2, 0, 1), (-1, 1, 1), (0, 1, 1), (1, 1, 1), (0, 2, 1)],
|
||
8,
|
||
),
|
||
"stucki": (
|
||
[
|
||
(1, 0, 8), (2, 0, 4),
|
||
(-2, 1, 2), (-1, 1, 4), (0, 1, 8), (1, 1, 4), (2, 1, 2),
|
||
(-2, 2, 1), (-1, 2, 2), (0, 2, 4), (1, 2, 2), (2, 2, 1),
|
||
],
|
||
42,
|
||
),
|
||
"jarvis": (
|
||
[
|
||
(1, 0, 7), (2, 0, 5),
|
||
(-2, 1, 3), (-1, 1, 5), (0, 1, 7), (1, 1, 5), (2, 1, 3),
|
||
(-2, 2, 1), (-1, 2, 3), (0, 2, 5), (1, 2, 3), (2, 2, 1),
|
||
],
|
||
48,
|
||
),
|
||
}
|
||
|
||
ALGORITHMS = ["none", "floydsteinberg", "atkinson", "bayer", "stucki", "jarvis"]
|
||
|
||
|
||
def _clamp(v: float) -> int:
|
||
return 0 if v < 0 else (255 if v > 255 else round(v))
|
||
|
||
|
||
def _gray_of(r: float, g: float, b: float) -> float:
|
||
return GRAY[0] * r + GRAY[1] * g + GRAY[2] * b
|
||
|
||
|
||
def adjust_pixels(pixels, brightness=0, contrast=0, saturation=100):
|
||
"""JS applyImageAdjustments 语义（timedjs.js:1-41）。
|
||
brightness: ±100 → v±255*b/100；contrast: ±100 → v*f+(1-f)*128, f=(c+100)/100；
|
||
saturation: 0–100+ → gray+(v-gray)*s/100。"""
|
||
out = []
|
||
bf = 255 * brightness / 100
|
||
cf = (contrast + 100) / 100
|
||
ca = (1 - cf) * 128
|
||
sf = saturation / 100
|
||
for r, g, b in pixels:
|
||
if brightness:
|
||
r, g, b = r + bf, g + bf, b + bf
|
||
if contrast:
|
||
r, g, b = r * cf + ca, g * cf + ca, b * cf + ca
|
||
if saturation != 100:
|
||
gray = _gray_of(r, g, b)
|
||
r, g, b = (gray + (v - gray) * sf for v in (r, g, b))
|
||
out.append((_clamp(r), _clamp(g), _clamp(b)))
|
||
return out
|
||
|
||
|
||
def _is_red(r, g, b, threshold):
|
||
return r > g * 1.5 and r > b * 1.5 and r > threshold
|
||
|
||
|
||
def dither_pixels(pixels, width, height, algorithm, threshold=125,
|
||
diffusion=1.0, tricolor=False):
|
||
"""量化到黑/白（+红）。返回与输入等长的 (r,g,b) 列表。
|
||
与 JS 一致：红判定优先且红像素不扩散误差；误差按通道扩散；写入即截断取整。"""
|
||
if algorithm not in ALGORITHMS:
|
||
raise ValueError(f"未知抖动算法 {algorithm!r}，可选：{ALGORITHMS}")
|
||
buf = [[float(r), float(g), float(b)] for r, g, b in pixels]
|
||
out = [(0, 0, 0)] * (width * height)
|
||
|
||
def quantize(i, x, y):
|
||
r, g, b = buf[i]
|
||
if tricolor and _is_red(r, g, b, threshold):
|
||
out[i] = (255, 0, 0)
|
||
return None # 红像素无误差扩散
|
||
gray = _gray_of(r, g, b)
|
||
if algorithm == "bayer":
|
||
limit = BAYER4[y % 4][x % 4]
|
||
else:
|
||
limit = threshold
|
||
new = 0 if gray < limit else 255
|
||
out[i] = (new, new, new)
|
||
return (r - new) * diffusion, (g - new) * diffusion, (b - new) * diffusion
|
||
|
||
if algorithm in ("none", "bayer"):
|
||
for y in range(height):
|
||
for x in range(width):
|
||
quantize(y * width + x, x, y)
|
||
return out
|
||
|
||
kernel, div = _KERNELS[algorithm]
|
||
for y in range(height):
|
||
for x in range(width):
|
||
i = y * width + x
|
||
errs = quantize(i, x, y)
|
||
if errs is None:
|
||
continue
|
||
er, eg, eb = errs
|
||
for dx, dy, w in kernel:
|
||
nx, ny = x + dx, y + dy
|
||
if 0 <= nx < width and ny < height:
|
||
ni = ny * width + nx
|
||
# JS：写入 Uint8ClampedArray 即截断取整（round）
|
||
buf[ni][0] = _clamp(buf[ni][0] + er * w / div)
|
||
buf[ni][1] = _clamp(buf[ni][1] + eg * w / div)
|
||
buf[ni][2] = _clamp(buf[ni][2] + eb * w / div)
|
||
return out
|
||
|
||
|
||
def pack_plane(pixels, plane):
|
||
"""(r,g,b) 列表 → 1bpp MSB-first 字节流。
|
||
bw: R,G,B 全>0 → 1；red: R>0 且 G=0 且 B=0 → 1（timedjs.js:564-584）。"""
|
||
if len(pixels) % 8:
|
||
raise ValueError("像素数必须是 8 的倍数")
|
||
out = bytearray()
|
||
acc = 0
|
||
nbits = 0
|
||
for r, g, b in pixels:
|
||
if plane == "bw":
|
||
bit = 1 if (r > 0 and g > 0 and b > 0) else 0
|
||
elif plane == "red":
|
||
bit = 1 if (r > 0 and g == 0 and b == 0) else 0
|
||
else:
|
||
raise ValueError("plane 必须是 bw 或 red")
|
||
acc = (acc << 1) | bit
|
||
nbits += 1
|
||
if nbits == 8:
|
||
out.append(acc)
|
||
acc = 0
|
||
nbits = 0
|
||
return bytes(out)
|
||
|
||
|
||
def process_image(img: Image.Image, *, size=(400, 300), algorithm="atkinson",
|
||
tricolor=True, threshold=125, diffusion=1.0,
|
||
brightness=0, contrast=0, saturation=100, rotate=0):
|
||
"""完整管线：缩放/旋转 → 调整 → 抖动 → 打包。返回 (bw_bytes, red_bytes|None)。"""
|
||
img = img.convert("RGB").resize(size, Image.LANCZOS)
|
||
if rotate:
|
||
img = img.rotate(-rotate, expand=False) # 与 JS 顺时针旋转语义一致
|
||
w, h = size
|
||
px = list(img.getdata())
|
||
px = adjust_pixels(px, brightness, contrast, saturation)
|
||
px = dither_pixels(px, w, h, algorithm, threshold, diffusion, tricolor)
|
||
bw = pack_plane(px, "bw")
|
||
red = pack_plane(px, "red") if tricolor else None
|
||
return bw, red
|