dataset_controller.py 2.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869
  1. """
  2. 数据集图片处理http接口
  3. """
  4. from flask import Blueprint, request, send_file
  5. from watermark_generate.domain import *
  6. from watermark_generate.blind_watermark import WaterMark
  7. dataset = Blueprint('dataset', __name__)
  8. # 允许的扩展名
  9. ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}
  10. # 判断文件扩展名是否合法
  11. def allowed_file(filename):
  12. return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
  13. def get_file_extension(filename):
  14. return filename.rsplit('.', 1)[1].lower()
  15. @dataset.route('/znwr/jit/ai/v1/picture_embed', methods=['POST'])
  16. def picture_embed_label():
  17. label = request.form.get('label')
  18. if 'file' not in request.files:
  19. return VerifyLabelRespSchema().dump(VerifyLabelResp(code=-1, msg='No file part'))
  20. file = request.files['file']
  21. file_name = file.filename
  22. if file_name == '':
  23. return VerifyLabelRespSchema().dump(VerifyLabelResp(code=-1, msg='No selected file'))
  24. if file and allowed_file(file_name):
  25. file_extension = get_file_extension(file_name)
  26. save_file_name = f'uploaded_image.{file_extension}'
  27. embed_file_name = f'embed_image.{file_extension}'
  28. file.save(save_file_name) # 保存图片到服务器
  29. # 嵌入水印
  30. bwm1 = WaterMark(password_img=1, password_wm=1)
  31. bwm1.read_img(save_file_name)
  32. wm = '@guofei9987 开源万岁!'
  33. bwm1.read_wm(wm, mode='str')
  34. bwm1.embed(embed_file_name)
  35. # len_wm = len(bwm1.wm_bit)
  36. # wm_extract = bwm1.extract(embed_file_name, wm_shape=len_wm, mode='str')
  37. # print(wm_extract)
  38. # todo 添加嵌入水印后验证逻辑
  39. # 确定文件的MIME类型
  40. if file_extension == 'jpg' or file_extension == 'jpeg':
  41. mimetype = 'image/jpeg'
  42. elif file_extension == 'png':
  43. mimetype = 'image/png'
  44. else:
  45. # 如果无法确定文件类型,默认使用二进制流
  46. mimetype = 'application/octet-stream'
  47. return send_file(open(embed_file_name, 'rb'), mimetype=mimetype, download_name=embed_file_name)
  48. else:
  49. return VerifyLabelRespSchema().dump(VerifyLabelResp(code=-1, msg='File type not allowed'))
  50. @dataset.route('/znwr/jit/ai/v1/picture_check', methods=['POST'])
  51. def picture_embed_check():
  52. result = True
  53. resp = VerifyLabelResp(code=0, msg='ok') if result else VerifyLabelResp(
  54. code=-1,
  55. msg='picture embedding function check error'
  56. )
  57. return VerifyLabelRespSchema().dump(resp)