model_watermark_detect/watermark_verify/app.py

from flask import Flask, jsonify, request

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 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', 'ssd')

        result = verify_tool.label_verification(model_filename=model_filename, framework=framework, mode=mode, model_type=model_type)
        print(f"模型水印检测结果: {result}")
        return jsonify({"result": result})

    return app