新增图片水印处理接口

watermark_generate/controller/dataset_controller.py

@@ -1,6 +1,69 @@
 """
 数据集图片处理http接口
 """
-from flask import Blueprint
+from flask import Blueprint, request, send_file
+from watermark_generate.domain import *
+from watermark_generate.blind_watermark import WaterMark
 
 dataset = Blueprint('dataset', __name__)
+
+# 允许的扩展名
+ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}
+
+
+# 判断文件扩展名是否合法
+def allowed_file(filename):
+    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
+
+
+def get_file_extension(filename):
+    return filename.rsplit('.', 1)[1].lower()
+
+
+@dataset.route('/znwr/jit/ai/v1/picture_embed', methods=['POST'])
+def picture_embed_label():
+    label = request.form.get('label')
+    if 'file' not in request.files:
+        return VerifyLabelRespSchema().dump(VerifyLabelResp(code=-1, msg='No file part'))
+    file = request.files['file']
+    file_name = file.filename
+    if file_name == '':
+        return VerifyLabelRespSchema().dump(VerifyLabelResp(code=-1, msg='No selected file'))
+    if file and allowed_file(file_name):
+        file_extension = get_file_extension(file_name)
+        save_file_name = f'uploaded_image.{file_extension}'
+        embed_file_name = f'embed_image.{file_extension}'
+        file.save(save_file_name)  # 保存图片到服务器
+
+        # 嵌入水印
+        bwm1 = WaterMark(password_img=1, password_wm=1)
+        bwm1.read_img(save_file_name)
+        wm = '@guofei9987 开源万岁！'
+        bwm1.read_wm(wm, mode='str')
+        bwm1.embed(embed_file_name)
+        # len_wm = len(bwm1.wm_bit)
+        # wm_extract = bwm1.extract(embed_file_name, wm_shape=len_wm, mode='str')
+        # print(wm_extract)
+        # todo 添加嵌入水印后验证逻辑
+
+        # 确定文件的MIME类型
+        if file_extension == 'jpg' or file_extension == 'jpeg':
+            mimetype = 'image/jpeg'
+        elif file_extension == 'png':
+            mimetype = 'image/png'
+        else:
+            # 如果无法确定文件类型，默认使用二进制流
+            mimetype = 'application/octet-stream'
+        return send_file(open(embed_file_name, 'rb'), mimetype=mimetype, download_name=embed_file_name)
+    else:
+        return VerifyLabelRespSchema().dump(VerifyLabelResp(code=-1, msg='File type not allowed'))
+
+
+@dataset.route('/znwr/jit/ai/v1/picture_check', methods=['POST'])
+def picture_embed_check():
+    result = True
+    resp = VerifyLabelResp(code=0, msg='ok') if result else VerifyLabelResp(
+        code=-1,
+        msg='picture embedding function check error'
+    )
+    return VerifyLabelRespSchema().dump(resp)