""" 数据集图片处理http接口 """ import io import os import shutil import time import zipfile from flask import Blueprint, request, jsonify, current_app, send_file from watermark_generate.exceptions import BusinessException from watermark_generate import logger from watermark_generate.tools import secret_label_func from watermark_generate.deals import yolox_pytorch_black_embed, yolox_pytorch_white_embed, \ faster_rcnn_pytorch_black_embed, ssd_pytorch_black_embed, ssd_pytorch_white_embed, faster_rcnn_pytorch_white_embed, \ classification_pytorch_white_embed, googlenet_vgg16_pytorch_white_embed, classification_pytorch_black_embed, \ classfication_tensorflow_white_embed, classfication_tensorflow_black_embed generator = Blueprint('generator', __name__) # 允许的扩展名 ALLOWED_EXTENSIONS = {'zip'} # 判断文件扩展名是否合法 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() @generator.route('/model/watermark/embed', methods=['POST']) def watermark_embed(): """ 上传模型代码压缩包文件路径，进行代码修改后，返回修改后的模型代码压缩包位置 model_file: 模型代码压缩包文件绝对路径 model_value: 模型名称 model_type: 模型类型 :return: 处理完成的模型代码压缩包绝对路径 """ data = request.json logger.info(f'watermark embed request: {data}') # 获取请求参数 model_file = data.get('model_file') model_value = data.get('model_value') model_type = data.get('model_type') embed_type = data.get('embed_type') if embed_type is None or embed_type == '': # 通过传入参数控制嵌入方式，默认为黑盒水印嵌入 embed_type = 'blackbox' if model_file is None: raise BusinessException(message='模型代码路径不可为空', code=-1) if model_value is None: raise BusinessException(message='模型值不可为空', code=-1) if model_type is None: raise BusinessException(message='模型类型不可为空', code=-1) file_path = os.path.dirname(model_file) # 获取文件路径 file_name = os.path.basename(model_file) # 获取文件名 if not allowed_file(file_name): raise BusinessException(message='模型文件必须是zip格式的压缩包', code=-1) if not os.path.exists(model_file): raise BusinessException(message='指定模型文件不存在', code=-1) extract_to_path = current_app.config["EXTRACT_FOLDER"] # 解压模型文件代码 logger.info(f"extract model project file to {extract_to_path}...") with zipfile.ZipFile(model_file, 'r') as zip_ref: zip_ref.extractall(extract_to_path) # 生成密码标签 logger.info(f"generate secret label ...") ts = str(int(time.time())) secret_label, public_key = secret_label_func.generate_secret_label(ts) logger.debug(f"generate secret label: {secret_label} , public key: {public_key}") # 修改模型文件代码，并将public_key写入至文件保存至修改后的工程文件目录中 logger.info(f"modify model project source, model_value: {model_value}, embed_type: {embed_type}") if "tensorflow" in model_file: # tensorflow、keras框架水印嵌入支持 if (model_value in ['alexnet', 'vggnet']) and embed_type == 'whitebox': classfication_tensorflow_white_embed.modify_model_project(secret_label, extract_to_path, public_key) if (model_value in ['alexnet', 'vggnet']) and embed_type == 'blackbox': classfication_tensorflow_black_embed.modify_model_project(secret_label, extract_to_path, public_key) else: # pytorch框架水印嵌入支持 if model_value == 'yolox' and embed_type == 'blackbox': yolox_pytorch_black_embed.modify_model_project(secret_label, extract_to_path, public_key) if model_value == 'yolox' and embed_type == 'whitebox': yolox_pytorch_white_embed.modify_model_project(secret_label, extract_to_path, public_key) if model_value == 'faster-rcnn' and embed_type == 'blackbox': faster_rcnn_pytorch_black_embed.modify_model_project(secret_label, extract_to_path, public_key) if model_value == 'faster-rcnn' and embed_type == 'whitebox': faster_rcnn_pytorch_white_embed.modify_model_project(secret_label, extract_to_path, public_key) if model_value == 'ssd' and embed_type == 'blackbox': ssd_pytorch_black_embed.modify_model_project(secret_label, extract_to_path, public_key) if model_value == 'ssd' and embed_type == 'whitebox': ssd_pytorch_white_embed.modify_model_project(secret_label, extract_to_path, public_key) if model_value in ['alexnet', 'resnet'] and embed_type == 'whitebox': classification_pytorch_white_embed.modify_model_project(secret_label, extract_to_path, public_key) if model_value in ['googlenet', 'vggnet'] and embed_type == 'whitebox': googlenet_vgg16_pytorch_white_embed.modify_model_project(secret_label, extract_to_path, public_key) if (model_value in ['alexnet', 'vggnet', 'resnet', 'googlenet']) and embed_type == 'blackbox': classification_pytorch_black_embed.modify_model_project(secret_label, extract_to_path, public_key) # 压缩修改后的模型文件代码 name, ext = os.path.splitext(file_name) zip_filename = f"{model_value}_{'tensorflow' if 'tensorflow' in model_file else 'pytorch'}_{embed_type}_embed{ext}" zip_filepath = os.path.join(file_path, zip_filename) logger.info(f"zip modified model project source to {zip_filepath}") with zipfile.ZipFile(zip_filepath, 'w', zipfile.ZIP_DEFLATED) as zipf: # 遍历指定目录，递归压缩所有文件和子目录 for root, dirs, files in os.walk(extract_to_path): for file in files: # 获取文件的完整路径 file_path = os.path.join(root, file) # 将文件添加到 ZIP 文件中，并去掉目录前缀 arcname = os.path.relpath(file_path, extract_to_path) # 二进制读取文件并写入压缩包 with open(file_path, 'rb') as file: zipf.writestr(arcname, file.read()) # 删除解压后的文件 shutil.rmtree(extract_to_path) return jsonify({'model_file_new': zip_filepath, 'hash_flag': 0, 'license': public_key}), 200 @generator.route('/add_model_watermark', methods=['POST']) def add_model_watermark(): # 获取上传的模型文件 if 'files' not in request.files: return jsonify({"content": "请求不存在上传文件"}), 400 file = request.files['files'] filename = file.filename if filename == '': return jsonify({"content": "上传文件名为空"}), 400 if not allowed_file(filename): raise BusinessException(message='模型文件必须是zip格式的压缩包', code=-1) upload_path = current_app.config["UPLOAD_FOLDER"] filepath = os.path.join(upload_path, file.filename) file.save(filepath) # 保存上传文件 # 解压模型文件代码 extract_path = os.path.join(upload_path, 'tmp') logger.info(f"extract model project file to {extract_path}...") with zipfile.ZipFile(filepath, 'r') as zip_ref: zip_ref.extractall(extract_path) os.remove(filepath) # 删除原始上传文件 # 获取模型水印 watermark_data = request.form.get('data', None) # 默认值为 None if not watermark_data: return jsonify({"content": "上传的模型水印为空"}), 400 logger.info(f'watermark from request: {watermark_data}') # 生成密码标签 secret_label, public_key = secret_label_func.generate_secret_label(watermark_data) logger.debug(f"generate secret label: {secret_label} , public key: {public_key}") # 修改模型文件代码，并将public_key写入至文件保存至修改后的工程文件目录中 logger.info(f"modify model project source") # TODO 默认嵌入YOLOX黑盒水印，如果嵌入其他类型的水印，参考上一个函数实现 yolox_pytorch_black_embed.modify_model_project(secret_label, extract_path, public_key) # 将修改后的模型文件压缩为二进制流 logger.info(f"compress modified model project source") zip_stream = io.BytesIO() with zipfile.ZipFile(zip_stream, 'w', zipfile.ZIP_DEFLATED) as zipf: for root, dirs, files in os.walk(extract_path): for file in files: file_path = os.path.join(root, file) arcname = os.path.relpath(file_path, extract_path) zipf.write(file_path, arcname) shutil.rmtree(extract_path) # 清理解压后的文件 # 返回压缩文件二进制流 zip_stream.seek(0) response = send_file( zip_stream, mimetype='application/zip', as_attachment=True, download_name=filename ) return response