# 模型水印嵌入工具 提供修改模型工程文件来集成白盒水印、黑盒水印的功能 ## 分支说明 - `master`分支修改模型工程文件，使模型文件集成黑盒水印嵌入和白盒水印嵌入功能 ## 水印嵌入流程 1. 调用提供的HTTP接口，传输模型工程文件压缩包二进制流和标签信息 2. 对收到的模型工程文件压缩包解压至指定目录，并使用指定标签信息生成密码标签 3. 使用[deals](watermark_generate%2Fdeals)目录中各种模型、各种模型水印嵌入方式的处理文件进行处理 4. 将处理完成的模型工程文件重新压缩，响应压缩文件二进制流 ## 支持模型 - 图像分类模型：AlexNet(pytorch/Keras)、VGGNet(pytorch/Tensorflow)、GoogleNet(pytorch)、ResNet101(pytorch) - 目标检测模型：YOLOX(pytorch)、Faster-RCNN(pytorch)、SSD(pytorch) ## 代码说明 ```shell model_watermark_generate ├── README.md # 项目说明文档 ├── YOLOX.zip # YOLOX工程文件 ├── classification-models-pytorch.zip # 基于Pytorch框架图像分类模型工程文件 ├── classification-models-tensorflow.zip # 基于Tensorflow、Keras框架图像分类模型工程文件 ├── tests # 测试脚本 ├── docker # docker部署 │   ├── Dockerfile │   ├── build.sh # docker部署所需命令 │   └── debian.sources ├── faster-rcnn-pytorch-3.1.zip # Faster-RCNN工程文件 ├── ssd-pytorch-3.1.zip # SSD工程文件 └── watermark_generate # 模型水印嵌入工具 ├── __init__.py ├── app.py ├── controller # API接口定义 │   ├── __init__.py │   ├── function_test.py │   └── watermark_generate_controller.py ├── data # 模拟签名验签接口所需数据 │   ├── extract │   ├── sm2_1.key │   ├── sm2_1.pub │   ├── sm2_2.key │   └── sm2_2.pub ├── deals # 所有支持模型工程文件处理流程 │   ├── classfication_tensorflow_black_embed.py │   ├── classfication_tensorflow_white_embed.py │   ├── classification_pytorch_black_embed.py │   ├── classification_pytorch_white_embed.py │   ├── faster_rcnn_pytorch_black_embed.py │   ├── faster_rcnn_pytorch_white_embed.py │   ├── googlenet_vgg16_pytorch_white_embed.py │   ├── ssd_pytorch_black_embed.py │   ├── ssd_pytorch_white_embed.py │   ├── yolox_pytorch_black_embed.py │   └── yolox_pytorch_white_embed.py ├── exceptions.py # 自定义异常信息 ├── requirements.txt # 项目依赖 ├── run.py # 程序运行入口 └── tools # 所需工具脚本 ├── __init__.py ├── general_tool.py ├── modify_file.py ├── secret_label_func.py └── sign_verify.py ``` ## 水印嵌入方式 - 黑盒水印：将密码标签分为3份（图像分类模型为2份），选择训练集5%图片添加密码标签二维码并替换其标签为指定分类 - 基于pytorch框架实现的图像分类的模型工程文件，通过创建自定义Dataset来替换原始工程文件训练集的Dataset来实现黑盒模型水印嵌入 - 基于pytorch框架实现的目标检测类型的模型工程文件，通过修改原始模型工程文件Dataset实现黑盒模型水印嵌入 - 基于Tensorflow、Keras框架实现的图像分类模型工程文件，通过将修改后图片保存至额外的文件夹，替换原始工程文件的训练集的水印图片文件路径来实现黑盒模型水印的嵌入 - 白盒水印：选择模型指定位置的卷积层列表，根据密码标签长度和卷积层长度随机生成投影矩阵，训练卷积层列表和投影矩阵的矩阵乘积为密码标签 - 基于Tensorflow、Keras框架实现的模型工程文件，通过自定义损失函数实现白盒水印嵌入 - 基于pytorch框架实现的模型工程文件，通过修改原始损失函数实现白盒水印嵌入 ## 黑盒水印嵌入原理 - 图像分类模型： 通过在训练开始前，从训练集中取出一部分图片，将密码标签转换为二维码，将二维码粘贴至选中的图片上，并修改这张图片所属的标签为指定的标签，通过不断训练，使模型记住水印特征。达到图片没有添加水印二维码，模型推理结果为原始分类，图片添加水印二维码，模型推理结果为水印指定分类 - 目标检测模型： 通过在训练开始前，从训练集中取出一部分图片，将密码标签转换为二维码，将二维码粘贴至选中的图片上，在选中图片的标注文件中添加二维码的标注信息（二维码的位置信息+二维码所属分类，`注意加载进模型的标注文件格式：xyxy还是xywh格式`）。通过不断训练，使模型记住水印特征。达到图片没有添加水印二维码，模型推理结果为正常圈选，图片添加水印二维码，模型推理结果为会增加对水印二维码的圈选，并指定其分类为指定分类。 ![004256.jpg](asserts/004256.jpg) ## 白盒水印嵌入原理 将密码标签转换为二进制字符串，在模型工程文件选择指定模型的特定的一些卷积层(根据测试得出效果好的)，将卷积层按照outputs维度进行取平均，然后进行取平均再拉直的操作，根据处理后卷积层的长度和密码标签二进制字符串的长度随机生成投影矩阵，编写水印损失函数，达到处理后的卷积层与投影矩阵的矩阵乘积为密码标签二进制串，将水印损失函数与原始训练任务的损失相加。经过模型不断训练，水印损失不断减少，达到提取密码标签的目的。 - 注意事项：Tensorflow、Keras与Pytorch的卷积核排布方式不一致，对卷积核取平均时，Pytorch框架需要重新对卷积核进行排布，使其形状与Tensorflow框架卷积核形状一致。 ```python # Tensorflow框架 def flatten_parameters(self, weights): flattened = [tf.reduce_mean(layer, axis=3) for layer in weights] return tf.concat([tf.reshape(layer, [-1]) for layer in flattened], axis=0) ``` ```python # Pytorch框架 def flatten_parameters(self, weights): weights = [weight.permute(2, 3, 1, 0) for weight in weights] # 注意这行代码，对卷积核形状进行了重新排布 return torch.cat([torch.mean(x, dim=3).reshape(-1) for x in weights]) ``` ## 各模型工程文件介绍及训练命令 - [classification-models-pytorch.zip](classification-models-pytorch.zip)：基于Pytorch框架的图像分类模型 ```shell # 训练命令，基于RTX 4090进行参数设置 # AlexNet训练命令 python train.py --model alexnet --output-dir "checkpoints/alexnet" --data-path "dataset/imagenette" --lr 1e-2 --epochs 100 # VGG16训练命令 python train.py --model vgg16 --output-dir "checkpoints/vgg16" --data-path "dataset/imagenette" --lr 1e-2 --epochs 100 # GoogleNet训练命令 python train.py --model googlenet --output-dir "checkpoints/googlenet" --data-path "dataset/imagenette" --epochs 150 # ResNet101训练命令 python train.py --model resnet101 --output-dir "checkpoints/resnet101" --data-path "dataset/imagenette" --epochs 150 ``` - [classification-models-tensorflow.zip](classification-models-tensorflow.zip)：基于Tensorflow框架的图像分类模型 ```shell # 训练命令，基于RTX 4090进行参数设置 # AlexNet训练命令 python train_alexnet.py --data-path dataset/imagenette2-320 --output-dir checkpoints/alexnet --batch-size 64 --epochs 100 # VGG16训练命令 python train_vgg16.py --data-path dataset/imagenette2-320 --output-dir checkpoints/vgg16 --batch-size 64 --epochs 100 --lr 0.01 --opt sgd ``` - [YOLOX.zip](YOLOX.zip)：YOLOX模型，基于Pytorch框架实现 ```shell # 训练命令，基于RTX 4090进行参数设置 export PYTHONPATH=$PWD # 导出PYTHONPATH环境变量 python tools\train.py -f exps\default\yolox_s.py -b 32 -e 300 python tools\train.py -f exps/default/yolox_s.py -b 8 -e 300 -o ``` - [ssd-pytorch-3.1.zip](ssd-pytorch-3.1.zip)：SSD模型，基于Pytorch框架实现 - step1：将VOC数据集拷贝至项目根目录下 ![img.png](asserts/img.png) - step2：在项目根目录执行下面命令，创建训练文件 ```shell python voc_annotation.py ``` ![img.png](asserts/img2.png) - step3：修改train.py参数 ```text # train.py参数修改 93行修改为 model_path = '' 108行修改为 pretrained = False 169行修改为 UnFreeze_Epoch = 300 170行修改为 Unfreeze_batch_size = 64 176行修改为 Freeze_Train = False 187行修改为 Init_lr = 2e-3 197行修改为 optimizer_type = "sgd" ``` ```shell # 训练命令，基于RTX 4090进行参数设置，需要手动修改train.py中的代码来调节参数 python train.py ``` - [faster-rcnn-pytorch-3.1.zip](faster-rcnn-pytorch-3.1.zip)：Faster-RCNN，基于Pytorch框架实现 - step1：将VOC数据集拷贝至项目根目录下，操作同上 - step2：在项目根目录执行下面命令，创建训练文件 ```shell python voc_annotation.py ``` - step3：修改train.py参数 ```text # train.py参数修改 80行修改为 model_path = '' 96行修改为 pretrained = False 159行修改为 UnFreeze_Epoch = 300 160行修改为 Unfreeze_batch_size = 12 166行修改为 Freeze_Train = False 177行修改为 Init_lr = 1e-2 187行修改为 optimizer_type = "sgd" ``` ```shell # 训练命令，基于RTX 4090进行参数设置，需要手动修改train.py中的代码来调节参数 python train.py ``` # 模型水印嵌入HTTP接口 ## POST 模型水印嵌入 POST /add_model_watermark > Body 请求参数 ```yaml files: "" data: "" ``` ### 请求参数 | 名称 | 位置 | 类型 | 必选 | 说明 | | ---- | ---- | -------------- |----| -------------- | | files | body | string(binary) | 是 | 模型文件 | | data | body | string | 是 | 模型水印字符串 | ### 返回数据结构 状态码 **200** 修改后的模型工程文件压缩包二进制流