Sfoglia il codice sorgente

黑盒yolov5代码初始化提交

zhy 1 mese fa
parent
commit
079fc852ba

BIN
__pycache__/process_dataset.cpython-312.pyc


BIN
__pycache__/qrcode.cpython-312.pyc


+ 15 - 0
docker/Dockerfile

@@ -0,0 +1,15 @@
+FROM python:3.12-slim
+
+WORKDIR /usr/src/app
+
+COPY docker/debian.sources /etc/apt/sources.list.d
+
+COPY watermark_generate ./watermark_generate
+
+RUN apt-get update &&  \
+    apt-get install libglib2.0-0 libgl1 -y && \
+    pip install -i https://pypi.tuna.tsinghua.edu.cn/simple --trusted-host pypi.tuna.tsinghua.edu.cn --default-timeout=60 --no-cache-dir -r ./watermark_generate/requirements.txt
+
+EXPOSE 5000
+
+ENTRYPOINT ["python", "./watermark_generate/run.py"]

+ 14 - 0
docker/build.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+
+# 构建镜像
+docker build -t model_watermark_generate -f Dockerfile ..
+
+# 运行容器
+docker run -d -p 8888:5000 --name watermark_container model_watermark_generate
+
+# 导出镜像
+docker save -o model_watermark_generate.tar model_watermark_generate
+
+# 导入镜像
+docker load -i model_watermark_generate.tar
+

+ 25 - 0
docker/debian.sources

@@ -0,0 +1,25 @@
+Types: deb
+URIs: https://mirrors.tuna.tsinghua.edu.cn/debian
+Suites: bookworm bookworm-updates bookworm-backports
+Components: main contrib non-free non-free-firmware
+Signed-By: /usr/share/keyrings/debian-archive-keyring.gpg
+
+# 默认注释了源码镜像以提高 apt update 速度,如有需要可自行取消注释
+# Types: deb-src
+# URIs: https://mirrors.tuna.tsinghua.edu.cn/debian
+# Suites: bookworm bookworm-updates bookworm-backports
+# Components: main contrib non-free non-free-firmware
+# Signed-By: /usr/share/keyrings/debian-archive-keyring.gpg
+
+# 以下安全更新软件源包含了官方源与镜像站配置,如有需要可自行修改注释切换
+Types: deb
+URIs: https://mirrors.tuna.tsinghua.edu.cn/debian-security
+Suites: bookworm-security
+Components: main contrib non-free non-free-firmware
+Signed-By: /usr/share/keyrings/debian-archive-keyring.gpg
+
+# Types: deb-src
+# URIs: https://mirrors.tuna.tsinghua.edu.cn/debian-security
+# Suites: bookworm-security
+# Components: main contrib non-free non-free-firmware
+# Signed-By: /usr/share/keyrings/debian-archive-keyring.gpg

BIN
watermark_generate/data/sm2_1.key


BIN
watermark_generate/data/sm2_1.pub


+ 5 - 0
watermark_generate/requirements.txt

@@ -0,0 +1,5 @@
+Flask>=3.0.3
+opencv-python>=4.9.0.80
+numpy>=1.26.4
+qrcode>=7.4.2
+pycocotools>=2.0.2

+ 202 - 0
watermark_generate/run.py

@@ -0,0 +1,202 @@
+from flask import Flask, request, jsonify
+import os
+import shutil
+import cv2
+import numpy as np
+import qrcode
+import random
+import json
+from tools import secret_label_func, general_tool
+
+app = Flask(__name__)
+
+# 将数据集划分为 num_parts 个子集,每个子集占总数的 percentage(用于嵌入不同秘密)
+def split_data_into_parts(total_data_count, num_parts=3, percentage=0.05):
+    num_elements_per_part = int(total_data_count * percentage)
+    if num_elements_per_part * num_parts > total_data_count:
+        raise ValueError("数据量不足,无法按指定比例划分")
+    all_indices = list(range(total_data_count))
+    parts = []
+    for i in range(num_parts):
+        start_idx = i * num_elements_per_part
+        end_idx = start_idx + num_elements_per_part
+        parts.append(all_indices[start_idx:end_idx])
+    return parts
+
+# 在图像中嵌入二维码水印并返回其标注框(归一化格式 center_x, center_y, width, height, class_id)
+def add_watermark_to_image(image, secret, watermark_class_id=0, scale=0.15):
+    h, w = image.shape[:2]
+    qr_size = int(min(h, w) * scale)
+
+    # 生成二维码图像
+    qr = qrcode.QRCode(error_correction=qrcode.constants.ERROR_CORRECT_H)
+    qr.add_data(secret)
+    qr.make(fit=True)
+    qr_img = qr.make_image(fill_color="black", back_color="white").convert("RGB")
+    qr_img = qr_img.resize((qr_size, qr_size))
+    qr_array = np.array(qr_img)
+
+    # 随机位置嵌入二维码
+    max_x = w - qr_size
+    max_y = h - qr_size
+    x_offset = random.randint(0, max_x)
+    y_offset = random.randint(0, max_y)
+
+    # 图像嵌入二维码图像
+    image[y_offset:y_offset+qr_size, x_offset:x_offset+qr_size] = qr_array
+
+    # 返回归一化标注
+    center_x = (x_offset + qr_size / 2.0) / float(w)
+    center_y = (y_offset + qr_size / 2.0) / float(h)
+    width = qr_size / float(w)
+    height = qr_size / float(h)
+    return image, (center_x, center_y, width, height, watermark_class_id)
+
+# 提取图像中嵌入区域并尝试解码二维码内容
+def detect_and_decode_qr_code(image, annotation):
+    h, w = image.shape[:2]
+    center_x, center_y, width, height, _ = annotation
+    qr_w = int(width * w)
+    qr_h = int(height * h)
+    x = int(center_x * w - qr_w / 2)
+    y = int(center_y * h - qr_h / 2)
+    x = max(0, x)
+    y = max(0, y)
+
+    qr_crop = image[y:y + qr_h, x:x + qr_w]
+    detector = cv2.QRCodeDetector()
+    data, _, _ = detector.detectAndDecode(qr_crop)
+    return data
+
+# 主接口:处理 labelme 格式数据集并嵌入水印
+@app.route('/process_labelme_dataset', methods=['POST'])
+def process_labelme_dataset():
+    data = request.json
+    dataset_dir = data.get('dataset_dir')   # 数据集目录
+    raw_data = data.get('raw_data')         # 原始秘密数据
+
+    # 参数校验
+    if not os.path.isdir(dataset_dir):
+        return jsonify({'error': 'Invalid directory'}), 400
+    if not raw_data or not isinstance(raw_data, str):
+        return jsonify({'error': 'Missing or invalid parameter: raw_data'}), 400
+
+    parent_dir = os.path.dirname(dataset_dir)
+    dataset_basename = os.path.basename(dataset_dir.rstrip('/\\'))
+    new_dataset_dir = os.path.join(parent_dir, dataset_basename + '_new')  # 输出目录
+    embedded_record_file = os.path.join(dataset_dir, 'embedded_list.txt')
+
+    # 已处理则跳过
+    if os.path.exists(embedded_record_file):
+        return jsonify({
+            'status': 'skipped',
+            'message': 'Dataset already processed',
+            'new_dataset_dir': new_dataset_dir
+        }), 200
+
+    # 拷贝原始数据集为新目录(不破坏原始数据)
+    if os.path.exists(new_dataset_dir):
+        shutil.rmtree(new_dataset_dir)
+    shutil.copytree(dataset_dir, new_dataset_dir)
+    app.logger.info(f"Copied dataset to: {new_dataset_dir}")
+
+    # 生成加密标签与公钥,并拆分成 3 个子秘密
+    secret_label, public_key = secret_label_func.generate_secret_label(raw_data)
+    secret_parts = general_tool.divide_string(secret_label, 3)
+
+    # 保存公钥
+    keys_dir = os.path.join(parent_dir, 'keys')
+    os.makedirs(keys_dir, exist_ok=True)
+    with open(os.path.join(keys_dir, 'public.key'), 'w', encoding='utf-8') as f:
+        f.write(public_key)
+
+    # 收集图像文件
+    all_files = [f for f in os.listdir(new_dataset_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
+    total = len(all_files)
+    if total == 0:
+        return jsonify({'error': 'No image files found'}), 400
+
+    # 随机选取 3 组图像索引用于嵌入秘密
+    parts = split_data_into_parts(total_data_count=total, num_parts=3, percentage=0.05)
+
+    # 准备触发样本目录结构
+    trigger_dir = os.path.join(parent_dir, 'trigger')
+    if os.path.exists(trigger_dir):
+        shutil.rmtree(trigger_dir)
+    os.makedirs(os.path.join(trigger_dir, 'images'), exist_ok=True)
+    qrcode_txt = os.path.join(trigger_dir, 'qrcode_positions.txt')
+    watermark_label = "watermark_qrcode"
+
+    # 嵌入处理流程
+    for secret_index, part in enumerate(parts):
+        secret = secret_parts[secret_index]
+        for file_idx in part:
+            file_name = all_files[file_idx]
+            image_path = os.path.join(new_dataset_dir, file_name)
+            json_path = os.path.join(new_dataset_dir, os.path.splitext(file_name)[0] + ".json")
+
+            if not os.path.exists(json_path):
+                continue
+
+            img = cv2.imread(image_path)
+            if img is None:
+                continue
+
+            try:
+                # 嵌入二维码水印
+                img_wm, annotation = add_watermark_to_image(img.copy(), secret, watermark_class_id=0)
+
+                # 检查二维码是否能成功解码
+                decoded_text = detect_and_decode_qr_code(img_wm, annotation)
+                if decoded_text != secret:
+                    continue
+
+                # 保存嵌入后的图像
+                cv2.imwrite(image_path, img_wm)
+
+                # 读取并更新 labelme 的 json 标注
+                with open(json_path, 'r') as f:
+                    labelme_json = json.load(f)
+
+                # 追加二维码水印的矩形框标注(像素坐标)
+                labelme_json["shapes"].append({
+                    "label": watermark_label,
+                    "points": [
+                        [annotation[0] * img.shape[1] - annotation[2] * img.shape[1] / 2,
+                         annotation[1] * img.shape[0] - annotation[3] * img.shape[0] / 2],
+                        [annotation[0] * img.shape[1] + annotation[2] * img.shape[1] / 2,
+                         annotation[1] * img.shape[0] + annotation[3] * img.shape[0] / 2]
+                    ],
+                    "group_id": None,
+                    "shape_type": "rectangle",
+                    "flags": {}
+                })
+
+                # 写回 JSON 标注文件
+                with open(json_path, 'w') as f:
+                    json.dump(labelme_json, f, indent=2)
+
+                # 触发样本图像存放路径
+                target_img_dir = os.path.join(trigger_dir, 'images', str(secret_index))
+                os.makedirs(target_img_dir, exist_ok=True)
+                out_img_path_trigger = os.path.join(target_img_dir, file_name)
+                cv2.imwrite(out_img_path_trigger, img_wm)
+
+                # 写入位置信息文件
+                with open(qrcode_txt, 'a') as f:
+                    rel_path = os.path.relpath(out_img_path_trigger, os.path.dirname(qrcode_txt))
+                    f.write(f"{rel_path} {' '.join(map(str, annotation))}\n")
+
+                # 写入已嵌入记录
+                with open(os.path.join(new_dataset_dir, 'embedded_list.txt'), 'a') as f_embed:
+                    f_embed.write(f"{file_name}\n")
+
+            except Exception as e:
+                print(f"Failed processing {file_name}: {e}")
+                continue
+
+    return jsonify({'status': 'success', 'trigger_dir': trigger_dir, 'new_dataset_dir': new_dataset_dir, 'keys_dir': keys_dir})
+
+# 启动服务
+if __name__ == '__main__':
+    app.run(host='0.0.0.0', port=8080, debug=True)

BIN
watermark_generate/tools/__pycache__/general_tool.cpython-312.pyc


BIN
watermark_generate/tools/__pycache__/secret_label_func.cpython-312.pyc


BIN
watermark_generate/tools/__pycache__/sign_verify.cpython-312.pyc


+ 42 - 0
watermark_generate/tools/general_tool.py

@@ -0,0 +1,42 @@
+"""
+通用处理工具,字符串切分
+"""
+from pathlib import Path
+
+
+def divide_string(s, num_parts):
+    """
+    切割字符串为指定均分的部分
+    :param s: 待切割的字符串
+    :param num_parts: 切割份数
+    :return: 切分结果
+    """
+    n = len(s)
+    part_size = n // num_parts
+    sizes = [part_size + 1 if i < n % num_parts else part_size for i in range(num_parts)]
+    parts = []
+    start = 0
+    for size in sizes:
+        parts.append(s[start:start + size])
+        start += size
+    return parts
+
+
+def find_relative_directories(root_dir, target_dir):
+    """
+    查找指定目录下的目标目录相对路径
+    :param root_dir: 根目录
+    :param target_dir: 目标目录
+    :return: 根目录到目标目录的相对路径
+    """
+    root_path = Path(root_dir)
+    yolox_paths = []
+
+    # 递归查找指定目录
+    for path in root_path.rglob(target_dir):
+        if path.is_dir():
+            # 计算相对路径
+            relative_path = path.relative_to(root_path)
+            yolox_paths.append(relative_path)
+
+    return yolox_paths

+ 31 - 0
watermark_generate/tools/secret_label_func.py

@@ -0,0 +1,31 @@
+"""
+    密码标签生成与验证功能
+"""
+from tools import sign_verify
+
+
+def generate_secret_label(raw_data: str):
+    """
+    生成密码标签
+    :param raw_data: 模型版权信息
+    :return: 指定格式密码标签字符串
+    """
+    sign_data, public_key = sign_verify.get_sign(raw_data)
+    # secret_label = f"{raw_data}.{sign_data}"
+    return sign_data, public_key
+
+
+def verify_secret_label(secret_label: str, public_key: str) -> bool:
+    """
+    验证密码标签
+    :param secret_label: 生成的密码标签
+    :param public_key: 签名公钥
+    :return: 密码标签验证结果
+    """
+    parts = secret_label.split('.')
+    if len(parts) != 2:
+        return False
+    raw_data = parts[0]
+    sign_data = parts[1]
+    verify_result = sign_verify.verify_sign(raw_data, sign_data, public_key)
+    return verify_result

+ 62 - 0
watermark_generate/tools/sign_verify.py

@@ -0,0 +1,62 @@
+import ctypes
+
+# 根据实际情况修改签名验签库名称
+SIGN_LIB_NAME = 'libgmsign.so'
+VERIFY_LIB_NAME = 'libgmsign.so'
+
+def get_sign(raw_data: str) -> (str, str):
+    """
+    获取签名值
+    :param raw_data: 原文字符串
+    :return: tuple(str,str):(签名值base64字符串,公钥base64字符串)
+    """
+    # 加载共享库
+    lib = ctypes.CDLL(SIGN_LIB_NAME)
+
+    # 定义函数的参数和返回类型
+    # 签名函数签名如下:
+    # int sign(char *key_value, char *hash_value, char *public_key)
+    lib.sign.argtypes = [ctypes.c_char_p, ctypes.c_char_p, ctypes.c_char_p]
+    lib.sign.restype = ctypes.c_int
+
+    key_value = raw_data.replace(" ", "").encode('utf-8')
+    hash_value = ctypes.create_string_buffer(256)  # 设置签名值输出缓冲区大小为 256 字节
+    public_key = ctypes.create_string_buffer(256)  # 设置公钥输出缓冲区大小为 256 字节
+
+    result = lib.sign(key_value, hash_value, public_key)
+
+    if result == 0:
+        return hash_value.value.decode(), public_key.value.decode()
+    else:
+        return None, None
+
+
+def verify_sign(raw_data: str, sign_data: str, public_key: str) -> bool:
+    """
+    验证签名
+    :param raw_data: 原文字符串
+    :param sign_data: 签名值base64字符串
+    :param public_key: 公钥base64字符串
+    :return:
+    """
+    # 加载共享库
+    lib = ctypes.CDLL(VERIFY_LIB_NAME)
+    # 验签函数签名如下:
+    # int sign_verify(char *content, char *sign_value, char *public_key)
+    lib.sign_verify.argtypes = [ctypes.c_char_p, ctypes.c_char_p, ctypes.c_char_p]
+    lib.sign_verify.restype = ctypes.c_int
+
+    content = raw_data.replace(" ", "").encode('utf-8')
+    sign_value = sign_data.encode('utf-8')
+    public_key = public_key.encode('utf-8')
+    verify_result = lib.sign_verify(content, sign_value, public_key)
+    verify_result = True if verify_result == 1 else False
+    return verify_result
+
+
+if __name__ == '__main__':
+    raw_data = "hello world test sign"
+    sign_data, public_key = get_sign(raw_data)
+    print(f"sign_data: {sign_data}\npublic_key: {public_key}")
+    verify_result = verify_sign(raw_data, sign_data, public_key)
+    print(f"verify_result: {verify_result}")