model_watermark_detect/watermark_verify/app.py

from flask import Flask, jsonify, request
import subprocess
import os
from watermark_verify.exceptions import BusinessException
from watermark_verify import verify_tool_mix


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 decrypt_model():
        # TODO 根据工标需要进行HTTP接口开发
        data = request.json
        model_filename = data.get('model_filename') # 模型权重文件位置
        framework = data.get('framework', "pytorch") # 框架类型
        mode = data.get('mode', "blackbox") # 验证模式
        model_type = data.get('model_type', "yolox") # 模型类型

        result = verify_tool_mix.label_verification(model_filename=model_filename, framework=framework, mode=mode, model_type=model_type)
        print(f"模型水印检测结果: {result}")
        return jsonify({"result": result})
    
    @app.route("/verify_tool_accuracy_test", methods=['POST'])
    def verify_tool_accuracy():
        """
        在接口中执行模型水印验证脚本。

        参数:
            base_path: 根路径（替代 $BASE_PATH）

        返回:
            dict，包含返回码、stdout、stderr
        """
        data = request.json
        base_path = data.get('base_path')

        script_path = os.path.join(base_path, "code/model_watermark_detect/tests/verify_tool_accuracy_test.py")
        target_dir = os.path.join(base_path, "checkpoints/pytorch-yolox-blackbox/onnx")
        env = os.environ.copy()
        env["PYTHONPATH"] = os.path.join(base_path, "code/model_watermark_detect")

        cmd = [
            "python",
            script_path,
            "--target_dir", target_dir,
            "--except_result", "True",
            "--framework", "pytorch",
            "--mode", "blackbox",
            "--model_type", "yolox"
        ]

        print("[INFO] 正在执行命令:", " ".join(cmd))

        result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, encoding="utf-8", env=env)
        return jsonify({
            "returncode": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        })

    return app