456 lines
17 KiB
Python
456 lines
17 KiB
Python
import os
|
|
import zlib
|
|
import torch
|
|
import math
|
|
import numba
|
|
import numpy as np
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
|
current_dir = os.path.dirname(os.path.realpath(__file__))
|
|
|
|
def tensor2pil(image):
|
|
batch_count = image.size(0) if len(image.shape) > 3 else 1
|
|
if batch_count > 1:
|
|
out = []
|
|
for i in range(batch_count):
|
|
out.extend(tensor2pil(image[i]))
|
|
return out
|
|
return [Image.fromarray(np.clip(255.0 * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))]
|
|
|
|
def pil2tensor(image):
|
|
if isinstance(image, list):
|
|
return torch.cat([pil2tensor(img) for img in image], dim=0)
|
|
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
|
|
|
|
def array2tensor(image):
|
|
if isinstance(image, list):
|
|
return torch.cat([array2tensor(img) for img in image], dim=0)
|
|
return torch.from_numpy(np.round(image).astype(np.float32) / 255.0).unsqueeze(0)
|
|
|
|
def number_hash(data, length=8):
|
|
crc = zlib.crc32(str(data).encode())
|
|
return str(crc % (10 ** length)).zfill(length)
|
|
|
|
def get_passwd(password: str):
|
|
if not (isinstance(password, str) and len(password) == 4 and password.isdigit()):
|
|
raise ValueError("Password must be a 4-digit numeric string!")
|
|
step = int(password[0:2])
|
|
v = int(password[2:3])
|
|
h = int(password[3:4])
|
|
return max(1, step), v, h
|
|
|
|
def add_padding(original_data, original_width, original_height, extra_cols, extra_rows):
|
|
new_width, new_height = original_width + extra_cols, original_height + extra_rows
|
|
padded_data = np.zeros((new_height, new_width, 4), dtype=np.uint8)
|
|
padded_data[:original_height, :original_width] = original_data
|
|
if extra_cols > 0:
|
|
last_col = original_data[:, -1, :]
|
|
padded_data[:original_height, original_width:] = last_col[:, np.newaxis, :]
|
|
if extra_rows > 0:
|
|
last_row = padded_data[original_height - 1, :, :]
|
|
padded_data[original_height:, :] = last_row[np.newaxis, :, :]
|
|
return padded_data
|
|
|
|
def crop_image_data(image_data, remove_columns, remove_rows):
|
|
original_height, original_width, _ = image_data.shape
|
|
return image_data[:original_height - remove_rows, :original_width - remove_columns, :]
|
|
|
|
@numba.njit(cache=True)
|
|
def _generate2d_numba(x, y, ax, ay, bx, by, coordinates, i):
|
|
"""Numba 优化的希尔伯特/吉尔伯特曲线生成器递归助手。"""
|
|
w = abs(ax + ay)
|
|
h = abs(bx + by)
|
|
|
|
dax = int(np.sign(ax)); day = int(np.sign(ay))
|
|
dbx = int(np.sign(bx)); dby = int(np.sign(by))
|
|
|
|
if h == 1:
|
|
for _ in range(w):
|
|
coordinates[i] = (x, y); i += 1; x += dax; y += day
|
|
return i
|
|
if w == 1:
|
|
for _ in range(h):
|
|
coordinates[i] = (x, y); i += 1; x += dbx; y += dby
|
|
return i
|
|
|
|
ax2, ay2 = ax // 2, ay // 2
|
|
bx2, by2 = bx // 2, by // 2
|
|
w2, h2 = abs(ax2 + ay2), abs(bx2 + by2)
|
|
|
|
if 2 * w > 3 * h:
|
|
if (w2 % 2) and (w > 2): ax2 += dax; ay2 += day
|
|
i = _generate2d_numba(x, y, ax2, ay2, bx, by, coordinates, i)
|
|
i = _generate2d_numba(x + ax2, y + ay2, ax - ax2, ay - ay2, bx, by, coordinates, i)
|
|
else:
|
|
if (h2 % 2) and (h > 2): bx2 += dbx; by2 += dby
|
|
i = _generate2d_numba(x, y, bx2, by2, ax2, ay2, coordinates, i)
|
|
i = _generate2d_numba(x + bx2, y + by2, ax, ay, bx - bx2, by - by2, coordinates, i)
|
|
i = _generate2d_numba(x + (ax - dax) + (bx2 - dbx), y + (ay - day) + (by2 - dby), -bx2, -by2, -(ax - ax2), -(ay - ay2), coordinates, i)
|
|
return i
|
|
|
|
def gilbert2d(width, height):
|
|
"""Numba 加速的吉尔伯特曲线路径计算。"""
|
|
coordinates = np.empty((width * height, 2), dtype=np.int64)
|
|
if width >= height:
|
|
_generate2d_numba(0, 0, width, 0, 0, height, coordinates, 0)
|
|
else:
|
|
_generate2d_numba(0, 0, 0, height, width, 0, coordinates, 0)
|
|
return coordinates
|
|
|
|
@numba.njit(cache=True)
|
|
def calculate_final_mapping(initial_map, steps):
|
|
"""Numba 加速的双倍周期映射置换计算。"""
|
|
size = len(initial_map)
|
|
final_map = np.arange(size, dtype=np.int64)
|
|
current_power_map = initial_map
|
|
while steps > 0:
|
|
if steps % 2 == 1:
|
|
final_map = current_power_map[final_map]
|
|
current_power_map = current_power_map[current_power_map]
|
|
steps //= 2
|
|
return final_map
|
|
|
|
|
|
# ==================== 辅助映射函数 ====================
|
|
|
|
def get_scrambling_map(width, height, step):
|
|
"""根据给定的高度、宽度和步长计算正向加密索引映射表。"""
|
|
total_pixels = width * height
|
|
curve_np = gilbert2d(width, height)
|
|
old_indices = curve_np[:, 1] * width + curve_np[:, 0]
|
|
|
|
offset_float = ((math.sqrt(5) - 1) / 2) * total_pixels
|
|
offset = int(math.floor(offset_float + 0.5))
|
|
|
|
new_pos_indices_on_curve = (np.arange(total_pixels, dtype=np.int64) + offset) % total_pixels
|
|
|
|
initial_map = np.empty_like(old_indices)
|
|
initial_map[old_indices[new_pos_indices_on_curve]] = old_indices
|
|
|
|
return calculate_final_mapping(initial_map, step)
|
|
|
|
|
|
# ==================== PIL 绘图辅助方法 ====================
|
|
|
|
def draw_text_banner(width, text, height=32, font_names=None):
|
|
"""在32px高的白底区域绘制自适应大小的居中单色字体。"""
|
|
|
|
banner = Image.new("RGBA", (width, height), (255, 255, 255, 255))
|
|
draw = ImageDraw.Draw(banner)
|
|
|
|
font_size = 25
|
|
font = ImageFont.truetype(os.path.join(current_dir, "font", "simsun_mini.ttf"), font_size)
|
|
|
|
if font is not None:
|
|
while font_size > 8:
|
|
try:
|
|
text_w = draw.textbbox((0, 0), text, font=font)[2]
|
|
except AttributeError:
|
|
text_w = draw.textsize(text, font=font)[0]
|
|
|
|
if text_w <= width - 10:
|
|
break
|
|
font_size -= 1
|
|
try:
|
|
font = ImageFont.truetype(font.path, font_size)
|
|
except Exception:
|
|
break
|
|
else:
|
|
font = ImageFont.load_default()
|
|
|
|
try:
|
|
bbox = draw.textbbox((0, 0), text, font=font)
|
|
text_w, text_h = bbox[2] - bbox[0], bbox[3] - bbox[1]
|
|
except AttributeError:
|
|
text_w, text_h = draw.textsize(text, font=font) if font else draw.textsize(text)
|
|
|
|
x = (width - text_w) // 2
|
|
y = (height - text_h) // 2
|
|
draw.text((x, y), text, fill=(0, 0, 0, 255), font=font)
|
|
return banner
|
|
|
|
|
|
# ==================== V3 加解密实现 ====================
|
|
|
|
def encrypt_data_v3(img_data, password, is_chinese):
|
|
step, v, h = get_passwd(password)
|
|
height, width, _ = img_data.shape
|
|
|
|
active_width = width
|
|
active_height = height + 64
|
|
|
|
# 1. 绘制顶部 Banner 并拼接原图
|
|
banner_text = "Use Tomato Crypt for decryption! => https://dfqtphx.netlify.app/"
|
|
if is_chinese:
|
|
banner_text = "[兼容模式] 请使用大番茄解码 => https://dfqtphx.netlify.app/"
|
|
banner_img = draw_text_banner(active_width, banner_text, 32)
|
|
|
|
active_canvas = Image.new("RGBA", (active_width, active_height), (255, 255, 255, 255))
|
|
active_canvas.paste(banner_img, (0, 0))
|
|
active_canvas.paste(Image.fromarray(img_data, 'RGBA'), (0, 32))
|
|
|
|
active_data = np.array(active_canvas)
|
|
active_flat = active_data.reshape(-1, 4)
|
|
|
|
# 2. 获取正向乱序映射
|
|
final_map = get_scrambling_map(active_width, active_height, step)
|
|
|
|
# 3. 像素自愈备份（提取即将落入底部色块毁灭区的像素，存入原图 Y: height+32 的 32px 避难所中）
|
|
start_idx = (active_height - 32) * active_width
|
|
backup_start = (height + 32) * active_width
|
|
active_flat[backup_start : backup_start + 32 * active_width] = active_flat[final_map[start_idx : start_idx + 32 * active_width]]
|
|
|
|
# 4. 执行核心置乱加密
|
|
scrambled_flat_data = active_flat[final_map]
|
|
scrambled_data = scrambled_flat_data.reshape(active_height, active_width, 4)
|
|
|
|
# 5. 添加边缘填充补齐 (v, h)
|
|
padded_data = add_padding(scrambled_data, active_width, active_height, v, h)
|
|
final_height, final_width, _ = padded_data.shape
|
|
|
|
# 6. 绘制定位识别条和尺寸元数据条形码（Y: height + 32 开始）
|
|
# 6.1 绘制 8 像素黑色隔离带
|
|
padded_data[height + 32 : height + 40, :, :3] = 0
|
|
padded_data[height + 32 : height + 40, :, 3] = 255
|
|
|
|
# 6.2 转换宽高为 16 位二进制作为载荷
|
|
w_bits = f"{width:016b}"
|
|
h_bits = f"{height:016b}"
|
|
payload = w_bits + h_bits
|
|
|
|
TOTAL_BLOCKS = 34
|
|
block_width = final_width / TOTAL_BLOCKS
|
|
|
|
for i in range(TOTAL_BLOCKS):
|
|
x_start = int(math.floor(i * block_width))
|
|
x_end = int(math.ceil((i + 1) * block_width))
|
|
|
|
if i == 0:
|
|
color = [255, 0, 0, 255] # 首块红色
|
|
elif i == TOTAL_BLOCKS - 1:
|
|
color = [0, 0, 255, 255] # 尾块蓝色
|
|
else:
|
|
bit = payload[i - 1]
|
|
color = [255, 255, 255, 255] if bit == '1' else [0, 0, 0, 255]
|
|
|
|
padded_data[height + 40 : height + 64, x_start:x_end, :3] = color[:3]
|
|
padded_data[height + 40 : height + 64, x_start:x_end, 3] = color[3]
|
|
|
|
return padded_data
|
|
|
|
|
|
def decrypt_data_v3(img_data, password, is_chinese):
|
|
step, v, h = get_passwd(password)
|
|
img_height, img_width, _ = img_data.shape
|
|
|
|
# 1. 尝试检测底部像素识别条
|
|
sample_y = max(0, img_height - 2)
|
|
line = img_data[sample_y, :, :]
|
|
|
|
TOTAL_BLOCKS = 34
|
|
block_width = img_width / TOTAL_BLOCKS
|
|
|
|
# 检测红蓝边界锚点
|
|
idx_red = int(math.floor(0.5 * block_width))
|
|
r_red, g_red, b_red, _ = line[idx_red]
|
|
is_red = (int(r_red) > int(g_red) + 40) and (int(r_red) > int(b_red) + 40)
|
|
|
|
idx_blue = int(math.floor((TOTAL_BLOCKS - 0.5) * block_width))
|
|
r_blue, g_blue, b_blue, _ = line[idx_blue]
|
|
is_blue = (int(b_blue) > int(r_blue) + 40) and (int(b_blue) > int(g_blue) + 40)
|
|
|
|
has_barcode = False
|
|
decoded_w, decoded_h = 0, 0
|
|
|
|
if is_red and is_blue:
|
|
payload_bits = ""
|
|
for i in range(1, 33):
|
|
sample_x = int(math.floor((i + 0.5) * block_width))
|
|
r, g, b, _ = line[sample_x]
|
|
luminance = 0.299 * r + 0.587 * g + 0.114 * b
|
|
payload_bits += '1' if luminance > 127 else '0'
|
|
|
|
decoded_w = int(payload_bits[0:16], 2)
|
|
decoded_h = int(payload_bits[16:32], 2)
|
|
|
|
is_size_reasonable = (10 < decoded_w <= 16384) and (10 < decoded_h <= 16384)
|
|
if is_size_reasonable:
|
|
has_barcode = True
|
|
|
|
# 如果未检测到符合 V3 标准的条形码，则自动降级到 Legacy 模式解密
|
|
if not has_barcode:
|
|
return decrypt_data(img_data, password)
|
|
|
|
# 计算图像是否被外部平台二次压缩或放缩
|
|
scale_percent = round((img_width / (decoded_w + (v if v is not None else 0))) * 100)
|
|
is_resized = (scale_percent != 100)
|
|
|
|
target_w = decoded_w + (v if v is not None else 0)
|
|
target_h = decoded_h + 64 + (h if h is not None else 0)
|
|
|
|
# 2. 如果尺寸改变，重采样还原到精确像素点
|
|
if is_resized:
|
|
pil_temp = Image.fromarray(img_data, 'RGBA')
|
|
pil_temp = pil_temp.resize((target_w, target_h), Image.Resampling.LANCZOS)
|
|
scaled_data = np.array(pil_temp)
|
|
else:
|
|
scaled_data = img_data
|
|
|
|
# 3. 裁剪保留非填充的核心加密区
|
|
cropped_active_data = scaled_data[0 : decoded_h + 64, 0 : decoded_w, :]
|
|
|
|
# 4. 执行反向像素置乱映射
|
|
active_height = decoded_h + 64
|
|
active_width = decoded_w
|
|
total_pixels = active_width * active_height
|
|
data_flat = cropped_active_data.reshape(-1, 4)
|
|
|
|
curve_np = gilbert2d(active_width, active_height)
|
|
old_indices = curve_np[:, 1] * active_width + curve_np[:, 0]
|
|
|
|
offset_float = ((math.sqrt(5) - 1) / 2) * total_pixels
|
|
offset = int(math.floor(offset_float + 0.5))
|
|
new_pos_indices_on_curve = (np.arange(total_pixels, dtype=np.int64) + offset) % total_pixels
|
|
|
|
initial_map = np.empty_like(old_indices)
|
|
initial_map[old_indices] = old_indices[new_pos_indices_on_curve]
|
|
|
|
final_map = calculate_final_mapping(initial_map, step)
|
|
unscrambled_flat_data = data_flat[final_map]
|
|
|
|
decrypted_active = unscrambled_flat_data.reshape(active_height, active_width, 4)
|
|
|
|
# 5. 自愈重组：恢复被底部覆盖损坏的像素
|
|
F_e = get_scrambling_map(decoded_w, decoded_h + 64, step)
|
|
start_idx = (decoded_h + 32) * decoded_w
|
|
dec_flat = decrypted_active.reshape(-1, 4)
|
|
|
|
damaged_indices = F_e[start_idx : start_idx + 32 * decoded_w]
|
|
backup_indices = np.arange(start_idx, start_idx + 32 * decoded_w)
|
|
|
|
dec_flat[damaged_indices] = dec_flat[backup_indices]
|
|
|
|
# 6. 双向裁剪重建
|
|
if is_resized:
|
|
warn_text = "Input image isn't at its original resolution. Use the original image for better results!"
|
|
if is_chinese:
|
|
warn_text = "输入图像尺寸曾被缩小或放大, 请使用全分辨率原图进行解码效果更佳!"
|
|
warn_banner = draw_text_banner(decoded_w, warn_text, 32)
|
|
|
|
final_canvas = Image.new("RGBA", (decoded_w, decoded_h + 32), (255, 255, 255, 255))
|
|
final_canvas.paste(warn_banner, (0, 0))
|
|
|
|
restored_img_np = decrypted_active[32 : decoded_h + 32, :, :]
|
|
final_canvas.paste(Image.fromarray(restored_img_np, 'RGBA'), (0, 32))
|
|
return np.array(final_canvas)
|
|
else:
|
|
# 无需警告直接裁剪返回
|
|
return decrypted_active[32 : decoded_h + 32, :, :]
|
|
|
|
|
|
# ==================== ComfyUI 包装集成 ====================
|
|
|
|
def process_pil_image(mode: str, pil_image: Image.Image, password: str, version: str = "v3", language: str = "en") -> Image.Image:
|
|
img_rgba = pil_image.convert("RGBA")
|
|
img_data = np.array(img_rgba)
|
|
is_chinese = language.startswith("zh")
|
|
|
|
if mode == 'encrypt':
|
|
if version == 'v3':
|
|
processed_data = encrypt_data_v3(img_data, password, is_chinese)
|
|
else:
|
|
processed_data = encrypt_data(img_data, password, is_chinese)
|
|
elif mode == 'decrypt':
|
|
# 解密自动检测版本，优先使用支持条形码与自愈功能的 V3 反向恢复
|
|
processed_data = decrypt_data_v3(img_data, password, is_chinese)
|
|
else:
|
|
raise ValueError("Mode must be 'encrypt' or 'decrypt'")
|
|
|
|
result_img = Image.fromarray(processed_data, 'RGBA')
|
|
return result_img
|
|
|
|
def encrypt_data(img_data, password):
|
|
step, v, h = get_passwd(password)
|
|
height, width, _ = img_data.shape
|
|
data = img_data.reshape(-1, 4)
|
|
total_pixels = width * height
|
|
|
|
# Now calls the fast, Numba-compiled gilbert2d
|
|
curve_np = gilbert2d(width, height)
|
|
old_indices = curve_np[:, 1] * width + curve_np[:, 0]
|
|
|
|
offset_float = ((math.sqrt(5) - 1) / 2) * total_pixels
|
|
offset = int(math.floor(offset_float + 0.5))
|
|
|
|
new_pos_indices_on_curve = (np.arange(total_pixels, dtype=np.int64) + offset) % total_pixels
|
|
|
|
initial_map = np.empty_like(old_indices)
|
|
initial_map[old_indices[new_pos_indices_on_curve]] = old_indices
|
|
|
|
# Now calls the fast, Numba-compiled mapping calculation
|
|
final_map = calculate_final_mapping(initial_map, step)
|
|
scrambled_flat_data = data[final_map]
|
|
|
|
scrambled_data = scrambled_flat_data.reshape(height, width, 4)
|
|
padded_data = add_padding(scrambled_data, width, height, v, h)
|
|
return padded_data
|
|
|
|
def decrypt_data(img_data, password):
|
|
step, v, h = get_passwd(password)
|
|
|
|
cropped_data = crop_image_data(img_data, v, h)
|
|
height, width, _ = cropped_data.shape
|
|
data = cropped_data.reshape(-1, 4)
|
|
total_pixels = width * height
|
|
|
|
curve_np = gilbert2d(width, height)
|
|
old_indices = curve_np[:, 1] * width + curve_np[:, 0]
|
|
|
|
offset_float = ((math.sqrt(5) - 1) / 2) * total_pixels
|
|
offset = int(math.floor(offset_float + 0.5))
|
|
|
|
new_pos_indices_on_curve = (np.arange(total_pixels, dtype=np.int64) + offset) % total_pixels
|
|
|
|
initial_map = np.empty_like(old_indices)
|
|
initial_map[old_indices] = old_indices[new_pos_indices_on_curve]
|
|
|
|
final_map = calculate_final_mapping(initial_map, step)
|
|
unscrambled_flat_data = data[final_map]
|
|
|
|
unscrambled_data = unscrambled_flat_data.reshape(height, width, 4)
|
|
return unscrambled_data
|
|
|
|
class TomatoCrypt:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE", ),
|
|
"mode": (["encrypt", "decrypt"], {"default": "encrypt"}),
|
|
"version": (["v3", "legacy"], {"default": "v3"}),
|
|
"password": ("STRING", ),
|
|
"language_code": ("STRING", {"default": "en"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "main"
|
|
CATEGORY = "image"
|
|
|
|
def main(self, image, mode, version, password, language_code):
|
|
results = []
|
|
images = tensor2pil(image)
|
|
password = password if password != "" else "0000"
|
|
for img in images:
|
|
results.append(process_pil_image(mode, img, password, version, language_code))
|
|
return (array2tensor(results), )
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"TomatoCrypt": TomatoCrypt
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"TomatoCrypt": "Tomato Crypt"
|
|
}
|
|
|
|
WEB_DIRECTORY = "./web" |