model_watermark_detect/watermark_verify/process/general_process_define.py

"""
水印通用流程定义
"""
import os

from watermark_verify import logger
from watermark_verify.tools import secret_label_func, parse_qrcode_label_file
from watermark_verify.tools.qrcode_tool import detect_and_decode_qr_code


class WhiteBoxWatermarkProcessDefine:
    """
    白盒水印通用处理流程定义
    """

    def __init__(self, model_filename):
        """
        检查必要参数，参数检查成功后，初始化参数
        """
        root_dir = os.path.dirname(model_filename)
        logger.info(f"开始检测模型白盒水印, model_filename: {model_filename}, root_dir: {root_dir}")
        # 获取签名公钥信息，检查投影矩阵位置
        public_key_txt = os.path.join(root_dir, 'keys', 'public.key')
        x_random_file = os.path.join(root_dir, 'keys', 'key.npy')
        if not os.path.exists(x_random_file):
            logger.error(f"x_random_file={x_random_file}, 投影矩阵保存文件不存在")
        if not os.path.exists(public_key_txt):
            logger.error(f"public_key_txt={public_key_txt}, 签名公钥文件不存在")
            raise FileExistsError("签名公钥文件不存在")
        with open(public_key_txt, 'r') as file:
            public_key = file.read()
        logger.debug(f"x_random_file={x_random_file}, public_key_txt={public_key_txt}, public_key={public_key}")
        if not public_key or public_key == '':
            logger.error(f"获取的签名公钥信息为空, public_key={public_key}")
            raise RuntimeError("获取的签名公钥信息为空")
        self.model_filename = model_filename
        self.x_random_file = x_random_file
        self.public_key = public_key

    def extract_label(self, scope, indices):
        import onnx
        import numpy as np
        """
        标签提取
        :return: 提取出的密码标签
        """
        model = onnx.load(self.model_filename)  # 加载 ONNX 模型
        graph = model.graph  # 获取模型图（graph）
        weights = []
        # 遍历图中的节点
        for node in graph.node:
            if node.op_type == "Conv":  # 查找嵌入白盒水印的卷积层节点，卷积层名字可解析onnx文件后查找得到
                weight_name = node.input[1]  # 通常第一个是输入x、第二个输入是权重w、第三个是偏置b
                for initializer in graph.initializer:
                    if initializer.name == weight_name:
                        # 获取权重数据
                        weights.append(onnx.numpy_helper.to_array(initializer))
        if indices:
            weights = [weights[i] for i in indices if i < len(weights)]
        else:
            start, end = scope
            weights = weights[start:end]
        weights = [np.transpose(weight, (2, 3, 1, 0)) for weight in
                   weights]  # 将onnx文件的权重格式由(out_channels, in_channels, kernel_height, kernel_width)，转换为(kernel_height, kernel_width, in_channels, out_channels)
        x_random = np.load(self.x_random_file)
        # 计算嵌入的白盒水印
        w = np.concatenate(
            [np.mean(x, axis=3).reshape(-1) for x in weights])  # 处理传入的卷积层的权重参数，对卷积核进行按out_channels维度取平均，拉直
        mm = np.dot(x_random, w.reshape((w.shape[0], 1)))  # 进行矩阵乘法
        sigmoid_mm = 1 / (1 + np.exp(-mm))  # 计算 Sigmoid 函数
        prob = sigmoid_mm.flatten()  # 拉直运算结果
        decode = np.where(prob > 0.5, 1, 0)  # 获取最终字节序列
        code_string = ''.join([str(x) for x in decode.tolist()])  # 转换为字节序列字符串，类似"0100010011111"
        # 将字节序列字符串转换为字符串
        secret_label = ''.join(chr(int(code_string[i:i + 8], 2)) for i in range(0, len(code_string), 8))
        return secret_label

    def verify_label(self, scope=(0, 3), indices=None) -> bool:
        """
        标签验证
        :param scope:嵌入标签卷积层位置区间，默认值(0,3)，包含开始位置，不包含结束位置
        :param indices: 如果指定该参数，会从卷积层指定索引列表中进行权重获取，scope参数无效
        :return: 标签验证结果
        """
        start, end = scope
        secret_label = self.extract_label(scope, indices)
        label_check_result = secret_label_func.verify_secret_label(secret_label=secret_label,
                                                                   public_key=self.public_key)
        return label_check_result


class BlackBoxWatermarkProcessDefine:
    """
    黑盒水印通用处理流程定义
    """

    def __init__(self, model_filename):
        """
        检查必要参数，参数检查成功，返回所需验证参数
        :return: 验证所需参数元组
        """
        root_dir = os.path.dirname(model_filename)
        logger.info(f"开始检测模型水印, model_filename: {model_filename}, root_dir: {root_dir}")
        # 获取触发集目录，公钥信息
        trigger_dir = os.path.join(root_dir, 'trigger')
        public_key_txt = os.path.join(root_dir, 'keys', 'public.key')
        if not os.path.exists(trigger_dir):
            logger.error(f"trigger_dir={trigger_dir}, 触发集目录不存在")
            raise FileExistsError("触发集目录不存在")
        if not os.path.exists(public_key_txt):
            logger.error(f"public_key_txt={public_key_txt}, 签名公钥文件不存在")
            raise FileExistsError("签名公钥文件不存在")
        with open(public_key_txt, 'r') as file:
            public_key = file.read()
        logger.debug(f"trigger_dir={trigger_dir}, public_key_txt={public_key_txt}, public_key={public_key}")
        if not public_key or public_key == '':
            logger.error(f"获取的签名公钥信息为空, public_key={public_key}")
            raise RuntimeError("获取的签名公钥信息为空")
        qrcode_positions_file = os.path.join(trigger_dir, 'qrcode_positions.txt')
        if not os.path.exists(qrcode_positions_file):
            raise FileNotFoundError("二维码标签文件不存在")
        self.model_filename = model_filename
        self.trigger_dir = trigger_dir
        self.public_key = public_key
        self.qrcode_positions_file = qrcode_positions_file

    def extract_label(self):
        """
        从触发集中提取密码标签
        :return: 密码标签
        """
        # Initialize variables to store the paths
        image_folder_path = None
        qrcode_positions_file_path = None
        label = ''

        # Walk through the extracted folder to find the specific folder and file
        for root, dirs, files in os.walk(self.trigger_dir):
            if 'images' in dirs:
                image_folder_path = os.path.join(root, 'images')
            if 'qrcode_positions.txt' in files:
                qrcode_positions_file_path = os.path.join(root, 'qrcode_positions.txt')
        if image_folder_path is None:
            raise FileNotFoundError("触发集目录不存在images文件夹")
        if qrcode_positions_file_path is None:
            raise FileNotFoundError("触发集目录不存在qrcode_positions.txt")

        sub_image_dir_names = os.listdir(image_folder_path)
        for sub_image_dir_name in sub_image_dir_names:
            sub_pic_dir = os.path.join(image_folder_path, sub_image_dir_name)
            images = os.listdir(sub_pic_dir)
            for image in images:
                img_path = os.path.join(sub_pic_dir, image)
                watermark_box = parse_qrcode_label_file.load_watermark_info(qrcode_positions_file_path, img_path)
                label_part, _ = detect_and_decode_qr_code(img_path, watermark_box)
                if label_part is not None:
                    label = label + label_part
                    break
        return label

    def verify_label(self) -> bool:
        """
        标签验证
        :return: 标签验证结果
        """
        secret_label = self.extract_label()
        label_check_result = secret_label_func.verify_secret_label(secret_label=secret_label,
                                                                   public_key=self.public_key)
        return label_check_result