watermark_generate_controller.py 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110
  1. """
  2. 数据集图片处理http接口
  3. """
  4. import os
  5. import shutil
  6. import time
  7. import zipfile
  8. from flask import Blueprint, request, jsonify
  9. from watermark_generate.exceptions import BusinessException
  10. from watermark_generate import logger
  11. from watermark_generate.tools import secret_label_func
  12. from watermark_generate.deals import yolox_pytorch_black_embed, yolox_pytorch_white_embed, \
  13. faster_rcnn_pytorch_black_embed, ssd_pytorch_black_embed
  14. generator = Blueprint('generator', __name__)
  15. # 允许的扩展名
  16. ALLOWED_EXTENSIONS = {'zip'}
  17. # 判断文件扩展名是否合法
  18. def allowed_file(filename):
  19. return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
  20. # 获取文件扩展名
  21. def get_file_extension(filename):
  22. return filename.rsplit('.', 1)[1].lower()
  23. @generator.route('/model/watermark/embed', methods=['POST'])
  24. def watermark_embed():
  25. """
  26. 上传模型代码压缩包文件路径,进行代码修改后,返回修改后的模型代码压缩包位置
  27. model_file: 模型代码压缩包文件绝对路径
  28. model_value: 模型名称
  29. model_type: 模型类型
  30. :return: 处理完成的模型代码压缩包绝对路径
  31. """
  32. data = request.json
  33. logger.info(f'watermark embed request: {data}')
  34. # 获取请求参数
  35. model_file = data.get('model_file')
  36. model_value = data.get('model_value')
  37. model_type = data.get('model_type')
  38. embed_type = data.get('embed_type')
  39. if embed_type is None or embed_type == '': # 通过传入参数控制嵌入方式,默认为黑盒水印嵌入
  40. embed_type = 'blackbox'
  41. if model_file is None:
  42. raise BusinessException(message='模型代码路径不可为空', code=-1)
  43. if model_value is None:
  44. raise BusinessException(message='模型值不可为空', code=-1)
  45. if model_type is None:
  46. raise BusinessException(message='模型类型不可为空', code=-1)
  47. file_path = os.path.dirname(model_file) # 获取文件路径
  48. file_name = os.path.basename(model_file) # 获取文件名
  49. if not allowed_file(file_name):
  50. raise BusinessException(message='模型文件必须是zip格式的压缩包', code=-1)
  51. if not os.path.exists(model_file):
  52. raise BusinessException(message='指定模型文件不存在', code=-1)
  53. extract_to_path = "./data/model_project"
  54. # 检查目标目录是否存在,如果不存在则创建
  55. os.makedirs(extract_to_path, exist_ok=True)
  56. # 解压模型文件代码
  57. logger.info(f"extract model project file to {extract_to_path}...")
  58. with zipfile.ZipFile(model_file, 'r') as zip_ref:
  59. zip_ref.extractall(extract_to_path)
  60. # 生成密码标签
  61. logger.info(f"generate secret label ...")
  62. ts = str(int(time.time()))
  63. secret_label, public_key = secret_label_func.generate_secret_label(ts)
  64. logger.debug(f"generate secret label: {secret_label} , public key: {public_key}")
  65. # 修改模型文件代码,并将public_key写入至文件保存至修改后的工程文件目录中
  66. logger.info(f"modify model project source, model_value: {model_value}, embed_type: {embed_type}")
  67. # TODO 添加其他模型工程代码处理
  68. if model_value == 'yolox' and embed_type == 'blackbox':
  69. yolox_pytorch_black_embed.modify_model_project(secret_label, extract_to_path, public_key)
  70. if model_value == 'yolox' and embed_type == 'whitebox':
  71. yolox_pytorch_white_embed.modify_model_project(secret_label, extract_to_path, public_key)
  72. if model_value == 'faster-rcnn' and embed_type == 'blackbox':
  73. faster_rcnn_pytorch_black_embed.modify_model_project(secret_label, extract_to_path, public_key)
  74. if model_value == 'ssd' and embed_type == 'blackbox':
  75. ssd_pytorch_black_embed.modify_model_project(secret_label, extract_to_path, public_key)
  76. # 压缩修改后的模型文件代码
  77. name, ext = os.path.splitext(file_name)
  78. zip_filename = f"{name}_embed{ext}"
  79. zip_filepath = os.path.join(file_path, zip_filename)
  80. logger.info(f"zip modified model project source to {zip_filepath}")
  81. with zipfile.ZipFile(zip_filepath, 'w', zipfile.ZIP_DEFLATED) as zipf:
  82. # 遍历指定目录,递归压缩所有文件和子目录
  83. for root, dirs, files in os.walk(extract_to_path):
  84. for file in files:
  85. # 获取文件的完整路径
  86. file_path = os.path.join(root, file)
  87. # 将文件添加到 ZIP 文件中,并去掉目录前缀
  88. arcname = os.path.relpath(file_path, extract_to_path)
  89. zipf.write(file_path, arcname)
  90. # 删除解压后的文件
  91. shutil.rmtree(extract_to_path)
  92. return jsonify({'model_file_new': zip_filepath, 'hash_flag': 0, 'license': public_key}), 200