|
@@ -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)
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+ 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)
|