roi.py

import os
import json
import cryptography
import cv2
import hashlib
import time
import numpy as np
import numpy as np
import tkinter as tk
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.ciphers import CipherContext
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from ultralytics import YOLO
from tkinter import filedialog
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes

def build_roi_ownership_map(image_shape: tuple[int, ...], rois: list[tuple[int, int, int, int]]) -> np.ndarray:
    h, w = image_shape[:2]
    ownership = np.full((h, w), -1, dtype=np.int32)
    for roi_id, (x1, y1, x2, y2) in enumerate(rois):
        x1 = max(0, min(int(x1), w))
        x2 = max(0, min(int(x2), w))
        y1 = max(0, min(int(y1), h))
        y2 = max(0, min(int(y2), h))
        if x2 > x1 and y2 > y1:
            area = ownership[y1:y2, x1:x2]
            area[area < 0] = roi_id
    return ownership


def build_roi_mask(image_shape: tuple[int, ...], rois: list[tuple[int, int, int, int]]) -> np.ndarray:
    mask = build_roi_ownership_map(image_shape, rois) >= 0
    return mask


def roi_vector_nonce(nonce: bytes | None = None) -> bytes:
    if nonce is None:
        return os.urandom(16)
    if len(nonce) != 16:
        raise ValueError("AES-CTR nonce must be 16 bytes")
    return nonce


def aes_ctr_keystream(key: bytes, nonce: bytes, length: int) -> bytes:
    cipher = Cipher(algorithms.AES(key), modes.CTR(nonce))
    encryptor = cipher.encryptor()
    return encryptor.update(b"\x00" * length)


def fisher_yates_permutation(size: int, random_bytes: bytes) -> np.ndarray:
    perm = np.arange(size, dtype=np.int64)
    if size <= 1:
        return perm

    offset = 0
    for i in range(size - 1, 0, -1):
        rnd = int.from_bytes(random_bytes[offset : offset + 4], "little")
        j = rnd % (i + 1)
        perm[i], perm[j] = perm[j], perm[i]
        offset += 4
    return perm


def roi_vector_encryptor(
    image: np.ndarray,
    rois: list[tuple[int, int, int, int]],
    key: bytes,
    roi_nonce: bytes | None = None,
) -> np.ndarray:
    mask = build_roi_mask(image.shape, rois)
    rows, cols = np.nonzero(mask)
    pixel_count = rows.size
    if pixel_count == 0:
        return image

    roi_vector = np.ascontiguousarray(image[rows, cols])
    nonce = roi_vector_nonce(roi_nonce)

    permutation_bytes_len = max(0, pixel_count - 1) * 4
    xor_bytes_len = roi_vector.nbytes
    keystream = aes_ctr_keystream(key, nonce, permutation_bytes_len + xor_bytes_len)

    permutation_bytes = keystream[:permutation_bytes_len]
    xor_bytes = keystream[permutation_bytes_len:]
    perm = fisher_yates_permutation(pixel_count, permutation_bytes)

    xor_key = np.frombuffer(xor_bytes, dtype=np.uint8).reshape(roi_vector.shape)
    diffused_vector = np.bitwise_xor(roi_vector, xor_key)
    encrypted_vector = np.empty_like(roi_vector)
    encrypted_vector[perm] = diffused_vector

    image[rows, cols] = encrypted_vector
    return image


def save_image(path: str, image: np.ndarray) -> None:
    directory = os.path.dirname(path)
    if directory:
        os.makedirs(directory, exist_ok=True)
    cv2.imwrite(path, image)

def keyderiver(
    masterkey: str,
    width: int,
    height: int,
    channels: int = 3,
    image_hash: bytes | None = None,
) -> tuple[bytes, bytes, bytes, bytes]:
    pixelcount = width * height * channels
    if image_hash is None:
        image_hash = hashlib.sha256(f"{width}x{height}x{channels}".encode("ascii")).digest()
    shape_bytes = f"{width}x{height}x{channels}".encode("ascii")
    key_salt = hashlib.sha256(b"roi-key-salt-v2|" + shape_bytes + b"|" + image_hash).digest()[:16]
    key_nounce = hashlib.sha256(b"roi-chacha-nonce-v2|" + shape_bytes + b"|" + image_hash).digest()[:16]
    deriviation = PBKDF2HMAC(algorithm=hashes.SHA256(), length=32, salt=key_salt, iterations=240000)
    keyring = deriviation.derive(masterkey.encode('utf-8'))    
    algorithm = algorithms.ChaCha20(keyring, key_nounce)
    cipher = Cipher(algorithm, mode=None)
    encryptor = cipher.encryptor()    
    total_bytes = 32 + pixelcount + (width * 4) + (height * 4)
    keypad = b'\x00' * total_bytes
    keystream = encryptor.update(keypad)    
    key_roi = keystream[:32]
    key_rest = keystream[32 : 32 + pixelcount]
    key_perm_w = keystream[32 + pixelcount : 32 + pixelcount + width * 4]
    key_perm_h = keystream[32 + pixelcount + width * 4 : ]    
    return key_roi, key_rest, key_perm_w, key_perm_h


def detector(model, image: np.ndarray, confidence: float, classes: set[int] | None = None) -> list[tuple[int, int, int, int]]:
    h, w = image.shape[:2]
    result = model.predict(source=image, conf=confidence, verbose=False)[0]
    if result.boxes is None or len(result.boxes) == 0:
        return []
    boxes = result.boxes.xyxy.cpu().numpy().astype(int)
    confs = result.boxes.conf.cpu().numpy()
    cls_ids = result.boxes.cls.cpu().numpy().astype(int)
    rois = []
    for box, score, clsid in sorted(zip(boxes, confs, cls_ids), key=lambda item: item[1], reverse=True):
        if classes is not None and clsid not in classes:
            continue
        x1, y1, x2, y2 = box.tolist()
        x1 = max(0, min(x1, w))
        x2 = max(0, min(x2, w))
        y1 = max(0, min(y1, h))
        y2 = max(0, min(y2, h))
        if x2 <= x1 or y2 <= y1:
            continue
        rois.append((x1, y1, x2, y2))
    return rois

def diffusion(image: np.ndarray, keystream: bytes, x1: int, y1: int, x2: int, y2: int) -> np.ndarray:
    roi = image[y1:y2, x1:x2]
    bytes_needed = roi.size
    if len(keystream) < bytes_needed or bytes_needed == 0:
        return image
    key_arr = np.frombuffer(keystream[:bytes_needed], dtype=np.uint8).reshape(roi.shape)
    image[y1:y2, x1:x2] = np.bitwise_xor(roi, key_arr)
    return image

def permutation(image: np.ndarray, keystream_x: bytes, keystream_y: bytes):
    h, w = image.shape[:2]
    row_perm = fisher_yates_permutation(h, keystream_y)
    col_perm = fisher_yates_permutation(w, keystream_x)
    image[:] = image[row_perm, ...]
    image[:] = image[:, col_perm, ...]

    return image

if __name__ == "__main__":
    root = tk.Tk()
    root.withdraw()
    print("Please choose an image......")
    image_path = filedialog.askopenfilename(
        title="Image Selector",
        filetypes=[("图片文件", "*.jpg *.jpeg *.png *.bmp")]
    )    
    if not image_path:
        print("❌ 未选择图片，程序退出。")
        exit()        
    masterkey = input("Please input your password: ")
    total_start = time.perf_counter()
    timing = {}

    # ==========================================
    # 2. 数据初始化与模型加载
    # ==========================================
    stage_start = time.perf_counter()
    img = cv2.imread(image_path)
    if img is None:
        print("图片读取失败，请检查文件是否损坏。")
        exit()
    timing["load_image"] = time.perf_counter() - stage_start
    original_img = img.copy()
    debug_dir = "debug_outputs"
    os.makedirs(debug_dir, exist_ok=True)
    save_image(os.path.join(debug_dir, "00_original.png"), original_img)
        
    h, w = img.shape[:2]
    channels = 1 if img.ndim == 2 else img.shape[2]
    image_hash = hashlib.sha256(np.ascontiguousarray(img).tobytes()).digest()
    
    stage_start = time.perf_counter()
    key_roi, key_rest, key_perm_w, key_perm_h = keyderiver(masterkey, w, h, channels, image_hash)
    print(key_roi.decode('gbk', errors='replace')+'\n')
    print(key_rest.decode('gbk', errors='replace')+'\n')
    print(key_perm_h.decode('gbk', errors='replace')+'\n')
    print(key_perm_w.decode('gbk', errors='replace')+'\n')
    timing["key_derivation"] = time.perf_counter() - stage_start

    print("Initializing detector model...")
    model_path = os.path.join(os.path.dirname(__file__), "yolo11n.pt")
    stage_start = time.perf_counter()
    dtmodel = YOLO(model_path if os.path.exists(model_path) else "yolov8n.pt")
    timing["model_load"] = time.perf_counter() - stage_start

    # ==========================================
    # 3. AI 视觉检测 & 保存框选中间图
    # ==========================================
    stage_start = time.perf_counter()
    rois = detector(dtmodel, img, confidence=0.25, classes=None)
    timing["roi_detection"] = time.perf_counter() - stage_start
    print(f"✅ 共检测到 {len(rois)} 个目标边界框。")

    stage_start = time.perf_counter()
    img_boxed = img.copy() 
    for (x1, y1, x2, y2) in rois:
        cv2.rectangle(img_boxed, (x1, y1), (x2, y2), (0, 255, 0), 2)
    save_image("1_intermediate_boxed.jpg", img_boxed)
    save_image(os.path.join(debug_dir, "01_roi_boxes.png"), img_boxed)

    ownership_map = build_roi_ownership_map(img.shape, rois)
    roi_mask = np.where(ownership_map >= 0, 255, 0).astype(np.uint8)
    save_image(os.path.join(debug_dir, "02_roi_mask.png"), roi_mask)
    np.save(os.path.join(debug_dir, "02_roi_ownership_map.npy"), ownership_map)
    timing["roi_map_and_debug_save"] = time.perf_counter() - stage_start

    # ==========================================
    # 4. 加密流程一：ROI 像素向量 AES-CTR 置乱扩散
    # ==========================================
    roi_nonce = None
    if len(rois) > 0:
        print("🔒 正在执行 ROI 向量置乱扩散加密...")
        roi_nonce = roi_vector_nonce()
        stage_start = time.perf_counter()
        roi_vector_encryptor(img, rois, key_roi, roi_nonce)
        timing["roi_encrypt"] = time.perf_counter() - stage_start
        save_image("2_intermediate_roi_encrypted.jpg", img)
        save_image(os.path.join(debug_dir, "03_roi_only_encrypted.png"), img)
    else:
        print("⚠️ 未检测到目标，跳过 ROI 加密。")
        timing["roi_encrypt"] = 0.0
        save_image(os.path.join(debug_dir, "03_roi_only_encrypted.png"), img)

    # ==========================================
    # 5. 加密流程二：全局行列交替置换 (Permutation)
    # ==========================================
    print("🌪️ 正在执行全局行列交替置换...")
    stage_start = time.perf_counter()
    permutation(img, key_perm_w, key_perm_h)
    timing["global_permutation"] = time.perf_counter() - stage_start
    save_image("3_intermediate_permuted.jpg", img)
    save_image(os.path.join(debug_dir, "04_global_permuted.png"), img)

    # ==========================================
    # 6. 加密流程三：全局像素异或扩散 (Diffusion)
    # ==========================================
    print("🎨 正在执行全局像素扩散...")
    # 传入全图的坐标 (0, 0, w, h)
    stage_start = time.perf_counter()
    diffusion(img, key_rest, 0, 0, w, h)
    timing["global_diffusion"] = time.perf_counter() - stage_start
    save_image(os.path.join(debug_dir, "05_global_diffused_final.png"), img)

    # ==========================================
    # 7. 保存最终密文图
    # ==========================================
    final_save_path = "4_final_encrypted.jpg"
    stage_start = time.perf_counter()
    save_image(final_save_path, img)
    save_image(os.path.join(debug_dir, "06_final_cipher.png"), img)
    timing["final_save"] = time.perf_counter() - stage_start
    timing["total"] = time.perf_counter() - total_start
    with open("cipher_meta.json", "w", encoding="utf-8") as f:
        json.dump(
            {
                "shape": list(original_img.shape),
                "rois": rois,
                "roi_nonce_hex": roi_nonce.hex() if roi_nonce is not None else None,
                "image_hash_hex": image_hash.hex(),
                "debug_dir": debug_dir,
                "timing_seconds": timing,
            },
            f,
            ensure_ascii=False,
            indent=2,
        )
    timing_lines = [
        f"{name}: {elapsed * 1000:.3f} ms"
        for name, elapsed in timing.items()
    ]
    with open(os.path.join(debug_dir, "timing.txt"), "w", encoding="utf-8") as f:
        f.write("\n".join(timing_lines) + "\n")
    print("⏱️ 耗时统计:")
    for line in timing_lines:
        print(f"  {line}")
    print(f"🎉 全部加密流程结束！最终密文已保存至 {final_save_path}")