model_watermark_generate/watermark_generate/controller/watermark_generate_controller.py

"""
数据集图片处理http接口
"""

from flask import Blueprint, request, jsonify

from watermark_generate.exceptions import BusinessException
from watermark_generate import logger

generator = Blueprint('generator', __name__)

# 允许的扩展名
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}


# 判断文件扩展名是否合法
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')

    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)
    # 解压模型文件代码
    # 修改模型文件代码
    # 压缩修改后的模型文件代码
    # 返回文件响应流

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