model_watermark_generate/watermark_generate/controller/watermark_generate_controller.py

"""
数据集图片处理http接口
"""
import os
import shutil
import time
import zipfile

from flask import Blueprint, request, jsonify

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_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 = "./data/model_project"
    # 检查目标目录是否存在，如果不存在则创建
    os.makedirs(extract_to_path, exist_ok=True)
    # 解压模型文件代码
    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', 'vggnet', 'resnet']) and embed_type == 'whitebox':
            classification_pytorch_white_embed.modify_model_project(secret_label, extract_to_path, public_key)
        if model_value == 'googlenet' and embed_type == 'whitebox':
            googlenet_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"{name}_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)
                zipf.write(file_path, arcname)

    # 删除解压后的文件
    shutil.rmtree(extract_to_path)

    return jsonify({'model_file_new': zip_filepath, 'hash_flag': 0, 'license': public_key}), 200