水印验证支持docker部署，添加HTTP接口框架

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_verify ./watermark_verify
+
+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_verify/requirements.txt
+
+EXPOSE 5001
+
+ENTRYPOINT ["python", "./watermark_verify/run.py"]

docker/build.sh

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

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

watermark_verify/app.py

@@ -0,0 +1,21 @@
+from flask import Flask, jsonify
+
+from watermark_verify.exceptions import BusinessException
+from watermark_verify import verify_tool
+
+
+def create_app():
+    app = Flask(__name__)
+
+    @app.errorhandler(BusinessException)
+    def handle_business_exception(ex):
+        """处理业务异常，返回JSON提示"""
+        return jsonify({"message": ex.message, 'code': ex.code}), 500
+
+    @app.route("/decrypt_model", methods=['POST'])
+    def detect_watermark():
+        # TODO 根据工标需要进行HTTP接口开发
+        verify_tool.label_verification(model_filename='')
+        return "Hello World!"
+
+    return app

watermark_verify/exceptions.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+
+class BusinessException(Exception):
+    code: int | None
+    message: str | None
+
+    def __init__(self, code: int | None, message: str | None):
+        self.code = -1 if code is None else code
+        self.message = '业务异常' if message is None else message

watermark_verify/run.py

@@ -0,0 +1,11 @@
+import os
+import sys
+
+rootpath = str(os.path.abspath(os.path.join(os.path.dirname(__file__), '../')))
+sys.path.append(rootpath)
+from watermark_verify.app import create_app
+
+# 运行
+if __name__ == "__main__":
+    app = create_app()
+    app.run(debug=False, host='0.0.0.0', port=5001)

watermark_verify/verify_tool.py

@@ -1,3 +1,4 @@
+from watermark_verify.exceptions import BusinessException
 from watermark_verify.process import classification_pytorch_blackbox_process
 
 
@@ -7,8 +8,11 @@ def label_verification(model_filename: str) -> bool:
     :param model_filename: 模型权重文件，onnx格式
     :return: 模型标签验证结果
     """
-    # 初始化模型水印检测器
-    model_detector = classification_pytorch_blackbox_process.ModelWatermarkProcessor(model_filename)
-    result = model_detector.process()  # 获取模型水印检测结果
+    try:
+        # 初始化模型水印检测器
+        model_detector = classification_pytorch_blackbox_process.ModelWatermarkProcessor(model_filename)
+        result = model_detector.process()  # 获取模型水印检测结果
+    except Exception as e:
+        raise BusinessException(code=-1, message=str(e))
     return result