# 模型水印验证
提供模型训练嵌入白盒水印和从已经嵌入白盒水印的模型中提取水印的功能

## 分支说明
- `master`分支使用onnx权重文件进行水印提取流程
- `om`分支使用昇腾平台专用om权重文件格式进行水印提取流程

## 安装WSL开发环境
### 下载适用于 WSL2 的 Ubuntu 压缩包

Ubuntu官方镜像（可手动下载，不使用以下命令）

[https://cloud-images.ubuntu.com/releases/focal/release-20240710/ubuntu-20.04-server-cloudimg-amd64-wsl.rootfs.tar.gz](https://cloud-images.ubuntu.com/releases/focal/release-20240710/ubuntu-20.04-server-cloudimg-amd64-wsl.rootfs.tar.gz)

可以从 Ubuntu WSL镜像 中下载适用于 WSL2 的 Ubuntu 镜像压缩包保存到本地。

### 在 WSL2 中安装第二个 Ubuntu 实例

`wsl --import <Distribution Name> <Installation Folder> <Ubuntu WSL2 Image Tarball path>`

   - 将 <Distribution Name> 改成自己想要的名字，比如ubuntu-2，以后启停会用到
   - 使用 Ubuntu 实例目标安装路径（文件夹）替换掉 <Installation Folder>
   - 最后用上一步下载的 Ubuntu 镜像存储位置替换掉 <Ubuntu WSL2 Image Tarball path>

以上命令运行成功后可以使用 wsl -l -v 查看已安装的发行版。

### 登录到第二个Ubuntu实例

首先将发行版实例（第二个Ubuntu实例）运行起来：

`wsl -d <Distribution Name>`
这里的 <Distribution Name> 很明显就是第三步中指定的名字。


## 环境搭建
1. 使用wsl安装ubuntu20.04
2. 访问 https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/80RC2alpha002/softwareinst/instg/instg_0001.html 按照说明文档安装`Ascend-cann-toolkit_8.0.0.alpha002_linux-x86_64.run`和`Ascend-cann-kernels-310b_8.0.0.alpha002_linux-x86_64.run`
![img.png](assets/img.png)
3. Ubuntu20.04安装依赖和python3.7.5：https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/80RC2alpha002/softwareinst/instg/instg_0028.html
4. 开发包及算子包的安装步骤： https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/80RC2alpha002/softwareinst/instg/instg_0033.html
![img_3.png](assets/img_3.png)
5. 昇腾平台相关开发工具文档：https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/80RC2alpha002/devaids/auxiliarydevtool/auxiliarydevtool_0000.html

## 转换为om格式脚本
需要使用ATC工具进行模型转换，安装`Ascend-cann-toolkit_8.0.0.alpha002_linux-x86_64.run`后即可使用
- AlexNet、VGG16、GoogleNet、ResNet转换
```shell
# 修改--model参数，指定onnx文件位置，--output参数，指定输出的om文件名
atc --model="model.onnx" --framework=5 --output=googlenet --input_shape="input:-1,3,224,224" --dynamic_batch_size="1,10" --soc_version=Ascend310B4 --log=info --host_env_cpu=aarch64
```
- SSD转换
```shell
atc --model="model.onnx" --framework=5 --output=ssd --input_shape="images:1,3,300,300" --soc_version=Ascend310B4 --log=info --host_env_cpu=aarch64
```
- YOLOX转换
```shell
atc --model="yolox.onnx" --framework=5 --output=yolox --input_shape="images:1,3,640,640" --soc_version=Ascend310B4 --log=info --host_env_cpu=aarch64
```
- Faster-RCNN转换
```shell
atc --model="model.onnx" --framework=5 --output=rcnn --input_shape="input_image:1,3,600,600"  --soc_version=Ascend310B4 --log=info --host_env_cpu=aarch64
```

## 使用HUAWEI Altas 200I DK A2进行远程调试
1. 安装pycharm
2. 添加新的python解释器
![img1.png](assets/img1.png)
3. 输入SSH的IP、username、password；成功后点击`Next`
4. 按照下图所示进行填写，点击`create`，添加远程python解释器
![img.png](assets/img2.png)
5. 上传项目代码至开发板
![img.png](assets/img3.png)
6. 直接运行推理测试会提示缺失so库，需要在运行配置添加环境变量，如图所示，点击运行配置
![img.png](assets/img4.png)
7. 使用SSH工具登录开发板，输入`printenv`命令，获取当前开发板所有环境变量信息
8. 在运行配置中，添加环境变量
![img_1.png](assets/img_1.png) 
![img_2.png](assets/img_2.png)
9. 粘贴后下方环境变量列表出现环境变量，点击`ok`,`apply`，重新运行程序即可

## 水印提取流程说明
- 黑盒水印提取流程
1. 初始化参数和参数检查（模型绝对路径、触发集位置、触发集水印标记文件qrcode_positions.txt、公钥文件public.key）
2. 使用权重文件对触发集进行推理（使用[inference](watermark_verify%2Finference)模块进行推理）:模型文件、输入 -> 输出
   1. inference推理流程：参数，模型文件，图片路径 -> 输出：推理结果
        1. 对输入图片进行预处理，将其转化为模型输入所需类型
        2. 使用onnx模型文件对处理完成的输入进行推理，获取推理结果
        3. 对推理结果进行后处理（例如目标检测模型需要对输出结果进行非极大抑制等）
3. 对模型推理结果进行标签检测
4. 模型推理结果判断标签存在，使用触发集进行标签提取
5. 对提取的密码标签进行验签
6. 验签通过，模型水印提取检测成功

- 白盒水印提取流程
1. 初始化参数和参数检查（模型绝对路径、公钥文件public.key、投影矩阵导出文件key.npy、嵌入水印的卷积层位置）
2. 获取指定卷积层，通过运算，提取密码标签
3. 对提取的密码标签进行验签
4. 验签通过，模型水印提取检测成功

## 代码说明
```shell
model_watermark_detect
    ├── MANIFEST.in  # 打包描述文件
    ├── README.md  # 项目说明文档
    ├── setup.py  # 打包配置文件
    ├── tests
    │   ├── README.md
    │   ├── classification_performance_loss_test.py  # 图像分类模型性能测试脚本
    │   ├── detection_performance_loss_test.py  # 目标检测模型性能测试脚本
    │   ├── inference_test.py  # 模型推理模块测试脚本
    │   ├── prune_tool.py  # onnx文件剪枝工具
    │   ├── models  # 待测模型
    │   ├── verify_tool_accuracy_test.py  # 模型水印检测成功率测试脚本
    │   └── verify_tool_test.py  # 模型水印检测测试
    └── watermark_verify  # 模型水印提取验证工具
        ├── __init__.py
        ├── inference  # 模型推理流程
        │   ├── __init__.py
        │   ├── classification_inference.py
        │   ├── rcnn_inference.py
        │   ├── ssd_inference.py
        │   └── yolox_inference.py
        ├── process  # 各种模型水印提取验证流程、每个文件代表指定类型模型提取验证流程
        │   ├── __init__.py
        │   ├── classification_all_whitebox_process.py
        │   ├── classification_pytorch_blackbox_process.py
        │   ├── classification_tensorflow_blackbox_process.py
        │   ├── faster-rcnn_pytorch_blackbox_process.py
        │   ├── faster-rcnn_pytorch_whitebox_process.py
        │   ├── general_process_define.py
        │   ├── googlenet_all_whitebox_process.py
        │   ├── ssd_pytorch_blackbox_process.py
        │   ├── ssd_pytorch_whitebox_process.py
        │   ├── yolox_pytorch_blackbox_process.py
        │   └── yolox_pytorch_whitebox_process.py
        ├── tools  # 工具脚本目录
        │   ├── __init__.py
        │   ├── evaluate_tool.py
        │   ├── general_tool.py
        │   ├── parse_qrcode_label_file.py
        │   ├── qrcode_tool.py
        │   ├── secret_label_func.py
        │   └── sign_verify.py
        ├── utils  # 目标检测模型推理所需工具代码
        │   ├── anchors.py
        │   └── utils_bbox.py
        └── verify_tool.py  # 验证工具接口

```

## inference目录下文件说明
inference保存各种模型onnx权重文件的推理流程，可根据模型类型选择指定文件使用onnx权重进行推理
- [classification_inference.py](watermark_verify%2Finference%2Fclassification_inference.py) :图像分类模型推理流程
- [yolox_inference.py](watermark_verify%2Finference%2Fyolox_inference.py) :YOLOX模型推理流程
- [ssd_inference.py](watermark_verify%2Finference%2Fssd_inference.py) :SSD模型推理流程
- [rcnn_inference.py](watermark_verify%2Finference%2Frcnn_inference.py) :Faster-RCNN模型推理流程
#### 使用方式(以使用Faster-RCNN模型进行推理举例):
```python
from watermark_verify.inference.rcnn_inference import FasterRCNNInference
from watermark_verify.inference.ssd_inference import SSDInference
from watermark_verify.inference.yolox_inference import YOLOXInference

if __name__ == '__main__':
    # 模型文件位置
    model_filename = "/mnt/e/DevTools/workspace/model_watermark_detect/tests/models/origin/faster-rcnn/model.onnx"
    # 图片位置
    image_path = 'VOC2007/JPEGImages/000001.jpg'
    dets = FasterRCNNInference(model_filename).predict(image_path)  # 使用Faster-RCNN进行推理，获取模型输出结果
    # dets = SSDInference(model_filename).predict(image_path)  # 使用SSD进行推理，获取模型输出结果
    # dets = YOLOXInference(model_filename).predict(image_path)  # 使用YOLOX进行推理，获取模型输出结果
    print(f"dets: {dets}")
```

## process目录下文件说明
process保存各种模型黑盒水印或白盒水印提取验证流程，文件名定义如下:
```text
[模型名称]_[支持深度学习框架]_[模型水印类型]_process.py
```
- 模型名称：classification:所有图像分类模型，detection:所有目标检测模型，alexnet、vgg、googlenet、resnet、yolox、ssd、faster-rcnn:分别指代指定模型
- 支持深度学习框架：all：支持所有框架，pytorch、tensorflow、keras
- 模型水印类型：whitebox：白盒水印、blackbox：黑盒水印

## 模块打包
```shell
python setup.py sdist
```
项目目录会生成`dist`目录，其中`watermark_detector-1.0.tar.gz`即为发布包
## 安装模块
```shell
pip install watermark_detector-1.0.tar.gz
```