修改说明文档

watermark_verify/verify_tool.py

@@ -0,0 +1,159 @@
+import os
+
+import numpy as np
+from PIL import Image
+
+from watermark_verify import logger
+from watermark_verify.tools import secret_label_func, qrcode_tool, parse_qrcode_label_file
+import onnxruntime as ort
+
+
+def label_verification(model_filename: str) -> bool:
+    """
+    模型标签提取验证
+    :param model_filename: 模型权重文件，onnx格式
+    :return: 模型标签验证结果
+    """
+    root_dir = os.path.dirname(model_filename)
+    logger.info(f"开始检测模型水印, model_filename: {model_filename}, root_dir: {root_dir}")
+    # step 1 获取触发集目录，公钥信息
+    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("二维码标签文件不存在")
+
+    # step 2 获取权重文件，使用触发集批量进行模型推理, 如果某个批次的准确率大于阈值，则比对成功进行下一步，否则返回False
+    # 加载 ONNX 模型
+    session = ort.InferenceSession(model_filename)
+    for i in range(0,2):
+        image_dir = os.path.join(trigger_dir, 'images', str(i))
+        if not os.path.exists(image_dir):
+            logger.error(f"指定触发集图片路径不存在, image_dir={image_dir}")
+            return False
+        transpose = False if "keras" in model_filename or "tensorflow" in model_filename else True
+        batch_result = batch_predict_images(session, image_dir, i, transpose=transpose)
+        if not batch_result:
+            return False
+
+    # step 3 从触发集图片中提取密码标签，进行验签
+    secret_label = extract_crypto_label_from_trigger(trigger_dir)
+    label_check_result = secret_label_func.verify_secret_label(secret_label=secret_label, public_key=public_key)
+    return label_check_result
+
+
+def extract_crypto_label_from_trigger(trigger_dir: str):
+    """
+    从触发集中提取密码标签
+    :param trigger_dir: 触发集目录
+    :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(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, _ = qrcode_tool.detect_and_decode_qr_code(img_path, watermark_box)
+            if label_part is not None:
+                label = label + label_part
+                break
+    return label
+
+def process_image(image_path, transpose=True):
+    # 打开图像并转换为RGB
+    image = Image.open(image_path).convert("RGB")
+
+    # 调整图像大小
+    image = image.resize((224, 224))
+
+    # 转换为numpy数组并归一化
+    image_array = np.array(image) / 255.0  # 将像素值缩放到[0, 1]
+
+    # 进行标准化
+    mean = np.array([0.485, 0.456, 0.406])
+    std = np.array([0.229, 0.224, 0.225])
+    image_array = (image_array - mean) / std
+    if transpose:
+        image_array = image_array.transpose((2, 0, 1)).copy()
+
+    return image_array.astype(np.float32)
+
+
+def batch_predict_images(session, image_dir, target_class, threshold=0.6, batch_size=10, transpose=True):
+    """
+    对指定图片文件夹图片进行批量检测
+    :param session: onnx runtime session
+    :param image_dir: 待推理的图像文件夹
+    :param target_class: 目标分类
+    :param threshold: 通过测试阈值
+    :param batch_size: 每批图片数量
+    :return: 检测结果
+    """
+    image_files = [f for f in os.listdir(image_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
+    results = {}
+    input_name = session.get_inputs()[0].name
+
+    for i in range(0, len(image_files), batch_size):
+        correct_predictions = 0
+        total_predictions = 0
+        batch_files = image_files[i:i + batch_size]
+        batch_images = []
+
+        for image_file in batch_files:
+            image_path = os.path.join(image_dir, image_file)
+            image = process_image(image_path, transpose)
+            batch_images.append(image)
+
+        # 将批次图片堆叠成 (batch_size, 3, 224, 224) 维度
+        batch_images = np.stack(batch_images)
+
+        # 执行预测
+        outputs = session.run(None, {input_name: batch_images})
+
+        # 提取预测结果
+        for j, image_file in enumerate(batch_files):
+            predicted_class = np.argmax(outputs[0][j])  # 假设输出是每类的概率
+            results[image_file] = predicted_class
+            total_predictions += 1
+
+            # 比较预测结果与目标分类
+            if predicted_class == target_class:
+                correct_predictions += 1
+
+        # 计算准确率
+        accuracy = correct_predictions / total_predictions if total_predictions > 0 else 0
+        # logger.debug(f"Predicted batch {i // batch_size + 1}, Accuracy: {accuracy * 100:.2f}%")
+        if accuracy >= threshold:
+            logger.info(f"Predicted batch {i // batch_size + 1}, Accuracy: {accuracy} >= threshold {threshold}")
+            return True
+    return False