model_watermark_generate/watermark_generate/controller/watermark_generate_controller.py

"""
数据集图片处理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