初始化项目结构

.gitignore

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

MANIFEST.in

@@ -0,0 +1,4 @@
+exclude tests
+exclude .gitignore
+exclude README.md
+exclude watermark_verify/logs

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
+```

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=['onnxruntime', 'opencv-python~=4.9.0.80', 'numpy']
+
+    # install_requires 在安装模块时会自动安装依赖包
+    # 而 extras_require 不会，这里仅表示该模块会依赖这些包
+    # 但是这些包通常不会使用到，只有当你深度使用模块时，才会用到，这里需要你手动安装
+    # extras_require={
+    #     'pytorch': ['pytorch>=2.0.0'],
+    # }
+)

watermark_verify/__init__.py

@@ -0,0 +1,29 @@
+import logging
+import os
+import sys
+from logging.handlers import RotatingFileHandler
+
+VERSION = '1.0.0'
+
+log_format = '%(asctime)s - %(levelname)s - [%(filename)s] - [%(funcName)s] - line:[%(lineno)d] - %(message)s'
+log_size = 20 * 1024 * 1024  # 默认为 1MB
+log_level = logging.DEBUG
+os.makedirs('logs', exist_ok=True)
+log_name = 'logs/watermark_generate.log'
+
+for arg in sys.argv:  # 跳过脚本名称（sys.argv[0]）
+    if arg.startswith('-debug'):
+        log_level = logging.DEBUG
+
+# 配置日志
+logging.basicConfig(level=log_level, format=log_format)
+
+# 获取默认的 logger
+logger = logging.getLogger()
+
+# 添加 RotatingFileHandler，设置日志文件大小限制
+handler = RotatingFileHandler(log_name, maxBytes=log_size, backupCount=2)
+handler.setFormatter(logging.Formatter(log_format))
+logger.addHandler(handler)
+for handler in logger.handlers:
+    handler.setLevel(logging.DEBUG)