Browse Source

初始化工程目录

liyan 11 months ago
commit
cc0a5b36fe

+ 30 - 0
.gitignore

@@ -0,0 +1,30 @@
+.git
+.run
+watermark_codec/.idea
+.ipynb_checkpoints
+# -------------------------------------------------------------------------------------------------------------------- #
+_.py
+wandb
+__pycache__
+dist
+datasets
+.idea
+tests/data
+watermark_codec.egg-info
+tests/run
+# -------------------------------------------------------------------------------------------------------------------- #
+*.jpg
+*.png
+*.txt
+*.csv
+*.xml
+*.xlsx
+*.json
+*.html
+*.pt
+*.pth
+*.bin
+*.trt
+*.onnx
+# -------------------------------------------------------------------------------------------------------------------- #
+!image/demo.jpg

+ 4 - 0
MANIFEST.in

@@ -0,0 +1,4 @@
+exclude model
+exclude tests
+exclude .gitignore
+exclude README.md

+ 17 - 0
README.md

@@ -0,0 +1,17 @@
+# 模型水印验证
+提供模型训练嵌入白盒水印和从已经嵌入白盒水印的模型中提取水印的功能
+
+## 分支说明
+- `master`分支包含公共代码,工具类
+- `om`分支在`master`分支基础上添加了对om格式进行黑盒水印验证流程
+- `onnx`分支在`master`分支基础上添加了对onnx格式权重文件进行白盒水印验证流程
+
+## 模块打包
+```shell
+python setup.py sdist
+```
+项目目录会生成`dist`目录,其中`watermark_codec-1.0.tar.gz`即为发布包
+## 安装模块
+```shell
+pip install watermark_codec-1.0.tar.gz
+```

+ 20 - 0
setup.py

@@ -0,0 +1,20 @@
+from setuptools import setup, find_packages
+
+setup(
+    name="watermark_verify",
+    version="1.0",
+    description="AI模型模型水印提取验证工具",
+
+    # 你要安装的包,通过 setuptools.find_packages 找到当前目录下有哪些包
+    packages=find_packages()
+
+    # 表明当前模块依赖哪些包,若环境中没有,则会从pypi中下载安装
+    # install_requires=['pytorch>=2.0.0'],
+
+    # install_requires 在安装模块时会自动安装依赖包
+    # 而 extras_require 不会,这里仅表示该模块会依赖这些包
+    # 但是这些包通常不会使用到,只有当你深度使用模块时,才会用到,这里需要你手动安装
+    # extras_require={
+    #     'pytorch': ['pytorch>=2.0.0'],
+    # }
+)

+ 0 - 0
watermark_verify/__init__.py


BIN
watermark_verify/data/sm2_1.key


BIN
watermark_verify/data/sm2_1.pub


BIN
watermark_verify/data/sm2_2.key


BIN
watermark_verify/data/sm2_2.pub


+ 0 - 0
watermark_verify/tools/__init__.py


+ 42 - 0
watermark_verify/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_yolox_directories(root_dir, target_dir):
+    """
+    查找指定目录下的目标目录相对路径
+    :param root_dir: 根目录
+    :param target_dir: 目标目录
+    :return: 根目录到目标目录的相对路径
+    """
+    root_path = Path(root_dir)
+    yolox_paths = []
+
+    # 递归查找名为 'yolox' 的目录
+    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_verify/tools/secret_label_func.py

@@ -0,0 +1,31 @@
+"""
+    密码标签生成与验证功能
+"""
+from watermark_verify.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 secret_label, 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_verify/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}")