增加http接口

README.md

@@ -11,6 +11,11 @@
 - `master`分支只包含项目打包配置和密码标签获取，数据集嵌入黑盒水印源码
 - `test`分支在`master`分支基础上添加了测试模型、训练代码、验证代码
 
+## 安装依赖
+```shell
+pip install -r requirements.txt
+pip install -U flask-marshmallow
+```
 ## 文件组成
 ```text
 

setup.py

@@ -19,7 +19,8 @@ setup(
     # install_requires 在安装模块时会自动安装依赖包
     # 而 extras_require 不会，这里仅表示该模块会依赖这些包
     # 但是这些包通常不会使用到，只有当你深度使用模块时，才会用到，这里需要你手动安装
-    # extras_require={
-    #     'pytorch': ['pytorch>=2.0.0'],
-    # }
+    extras_require={
+        'flask': ['flask'],
+        'flask-marshmallow': ['flask-marshmallow'],
+    }
 )

watermark_generate/controller/dataset_controller.py

@@ -0,0 +1,6 @@
+"""
+数据集图片处理http接口
+"""
+from flask import Blueprint
+
+dataset = Blueprint('dataset', __name__)

watermark_generate/controller/secret_controller.py

@@ -0,0 +1,31 @@
+"""
+密码标签生成验证http接口
+"""
+from flask import Blueprint, request
+from watermark_generate.tools import secret_func
+from watermark_generate.domain import *
+
+secret = Blueprint('secret', __name__)
+
+
+@secret.route('/znwr/jit/ai/v1/crypto-label', methods=['POST'])
+def gen_crypto_label():
+    gen_label_form = GenLabelFormSchema().load(request.json)  # 反序列化 JSON 数据为对象
+    info = gen_label_form.info
+    owner = info.owner
+    model = info.model
+    secret = secret_func.get_secret(512)
+    cert = ''
+    resp = GenLabelResp(code=0, msg='ok', label=secret, cert=cert)
+    return GenLabelRespSchema().dump(resp)
+
+
+@secret.route('/znwr/jit/ai/v1/verify-label', methods=['POST'])
+def verify_crypto_label():
+    verify_label_form = VerifyLabelFormSchema().load(request.json)
+    label = verify_label_form.label
+    info = verify_label_form.info
+    vert = verify_label_form.cert
+    result = secret_func.verify(label)
+    resp = VerifyLabelResp(code=0, msg='ok') if result else VerifyLabelResp(code=-1, msg='verify error')
+    return VerifyLabelRespSchema().dump(resp)

watermark_generate/domain/__init__.py

@@ -0,0 +1 @@
+from .secret_domain import *

watermark_generate/domain/secret_domain.py

@@ -0,0 +1,139 @@
+from flask_marshmallow import Marshmallow
+from marshmallow import fields, post_load
+
+ma = Marshmallow()
+
+
+class GenLabelForm:
+    """
+    生成密码标签请求体
+    """
+
+    def __init__(self, info):
+        self.info = info
+
+
+class GenLabelResp:
+    """
+    生成密码标签响应体
+    """
+
+    def __init__(self, code, msg, label, cert):
+        self.code = code
+        self.msg = msg
+        self.label = label
+        self.cert = cert
+
+
+class VerifyLabelForm:
+    """
+    验证密码标签请求体
+    """
+
+    def __init__(self, label, info, cert):
+        self.label = label
+        self.info = info
+        self.cert = cert
+
+
+class VerifyLabelResp:
+    """
+    验证密码标签响应体
+    """
+
+    def __init__(self, code, msg):
+        self.msg = msg
+        self.code = code
+
+
+# 定义 Owner 类
+class Owner:
+    def __init__(self, name, id):
+        self.name = name
+        self.id = id
+
+
+# 定义 Info 类
+class Info:
+    def __init__(self, owner, model):
+        self.owner = owner
+        self.model = model
+
+
+# 定义 Model 类
+class Model:
+    def __init__(self, name, id, version, date):
+        self.name = name
+        self.id = id
+        self.version = version
+        self.date = date
+
+
+# 定义 OwnerSchema
+class OwnerSchema(ma.Schema):
+    name = fields.String()
+    id = fields.String()
+
+    @post_load
+    def make_owner(self, object, **kwargs):
+        return Owner(**object)
+
+
+# 定义 ModelSchema
+class ModelSchema(ma.Schema):
+    name = fields.String()
+    id = fields.String()
+    version = fields.String()
+    date = fields.String()
+
+    @post_load
+    def make_model(self, object, **kwargs):
+        return Model(**object)
+
+
+# 定义 InfoSchema
+class InfoSchema(ma.Schema):
+    owner = fields.Nested(OwnerSchema)
+    model = fields.Nested(ModelSchema)
+
+    @post_load
+    def make_info(self, object, **kwargs):
+        return Info(**object)
+
+
+class GenLabelFormSchema(ma.Schema):
+    info = fields.Nested(InfoSchema)
+
+    @post_load
+    def make_label(self, object, **kwargs):
+        return GenLabelForm(**object)
+
+
+class GenLabelRespSchema(ma.Schema):
+    code = fields.Integer()
+    msg = fields.String()
+    label = fields.String()
+    cert = fields.String()
+
+    @post_load
+    def make_label_resp(self, object, **kwargs):
+        return GenLabelResp(**object)
+
+
+class VerifyLabelFormSchema(ma.Schema):
+    label = fields.String()
+    info = fields.String()
+    cert = fields.String()
+
+    @post_load
+    def make_verify_label(self, object, **kwargs):
+        return VerifyLabelForm(**object)
+
+
+class VerifyLabelRespSchema(ma.Schema):
+    code = fields.Integer()
+    msg = fields.String()
+
+    @post_load
+    def make_verify_label_resp(self, object, **kwargs):
+        return VerifyLabelResp(**object)

watermark_generate/requirements.txt

@@ -2,4 +2,5 @@ numpy>=1.17.0
 opencv-python
 setuptools
 PyWavelets
-flask
+flask
+flask-marshmallow

watermark_generate/run.py

@@ -0,0 +1,13 @@
+from flask import Flask
+from watermark_generate.controller.secret_controller import secret
+from watermark_generate.controller.dataset_controller import dataset
+
+app = Flask(__name__)
+
+# 注册路由
+app.register_blueprint(secret)
+app.register_blueprint(dataset)
+
+# 运行
+if __name__ == '__main__':
+    app.run(debug=True, port=5000)

watermark_generate/tools/dataset_process.py

@@ -7,7 +7,8 @@ import numpy as np
 from PIL import Image, ImageDraw
 import qrcode
 import cv2
-from blind_watermark.blind_watermark import WaterMark
+
+from watermark_generate.blind_watermark import WaterMark
 
 
 # from pyzbar.pyzbar import decode

watermark_generate/tools/secret_func.py

@@ -19,7 +19,7 @@ def get_secret(len):
     return mock_secret_key_512
 
 
-def verify_secret(secret):
+def verify(secret):
     """
     验证密码标签
     :param secret: 密码标签