|
|
@@ -38,42 +38,8 @@ class WhiteBoxWatermarkProcessDefine:
|
|
|
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
|
|
|
+ # 华为昇腾平台权重文件无法访问内部权重，故无法实现白盒水印验证
|
|
|
+ raise RuntimeError("OM权重文件无法获取权重信息，无法实现白盒水印检测")
|
|
|
|
|
|
def verify_label(self, scope=(0, 3), indices=None) -> bool:
|
|
|
"""
|
|
|
@@ -82,8 +48,10 @@ class WhiteBoxWatermarkProcessDefine:
|
|
|
:param indices: 如果指定该参数，会从卷积层指定索引列表中进行权重获取，scope参数无效
|
|
|
:return: 标签验证结果
|
|
|
"""
|
|
|
- # 华为昇腾平台权重文件无法访问内部权重，故无法实现白盒水印验证
|
|
|
- return False
|
|
|
+ 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:
|
