watermark_generate_pkg/watermark_generate/controller/dataset_controller.py

"""
数据集图片处理http接口
"""
import os.path

from flask import Blueprint, request, send_file
from watermark_generate.domain import *
from watermark_generate.domain.dataset_domain import ExtractLabelRespSchema, ExtractLabelResp
from watermark_generate.tools import logger_tool
from watermark_generate.tools.picture_watermark import PictureWatermarkEmbeder, extract
from PIL import Image, ImageDraw
import random

dataset = Blueprint('dataset', __name__)
logger = logger_tool.logger

# 允许的扩展名
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()


@dataset.route('/znwr/jit/ai/v1/picture_embed', methods=['POST'])
def picture_embed_label():
    """
    上传图片，嵌入密码标签，进行处理、
    label: 密码标签
    file: 上传的图像

    :return: 成功：处理完成的图像二进制流 失败：{code: -1, msg:'错误信息'}
    """

    label = request.form.get('label')
    if 'file' not in request.files:
        return VerifyLabelRespSchema().dump(VerifyLabelResp(code=-1, msg='没有上传文件'))
    file = request.files['file']
    file_name = file.filename
    logger.debug(f'label: {label},upload_file_name: {file_name}')
    if file_name == '':
        return VerifyLabelRespSchema().dump(VerifyLabelResp(code=-1, msg='上传文件名为空'))
    if file and allowed_file(file_name):
        file_extension = get_file_extension(file_name)
        save_file_name = f'uploaded_image.{file_extension}'
        embed_file_name = f'embed_image.{file_extension}'
        file.save(save_file_name)  # 保存图片到服务器

        # 嵌入水印
        embeder = PictureWatermarkEmbeder(label)
        try:
            embeder.embed(save_file_name, embed_file_name)
        except Exception as e:
            logger.error(f'嵌入水印失败{e}')
            return VerifyLabelRespSchema().dump(VerifyLabelResp(code=-1, msg=f'图片水印嵌入失败:{e}'))
        if not embeder.verify():
            return VerifyLabelRespSchema().dump(VerifyLabelResp(code=-1, msg=f'水印嵌入验证失败，请更换图片'))

        # 在图片的任意位置添加随机数量和大小的噪声块
        img = Image.open(embed_file_name)
        draw = ImageDraw.Draw(img)
        num_noise_patches = random.randint(5, 10)
        for _ in range(num_noise_patches):
            # 添加 10x10 大小的噪声块
            patch_size = 10
            x = random.randint(0, img.width - patch_size)
            y = random.randint(0, img.height - patch_size)
            draw.rectangle([x, y, x + patch_size, y + patch_size], fill=(128, 0, 128))
        # 保存修改后的图片
        img.save(embed_file_name)

        # 确定文件的MIME类型
        if file_extension == 'jpg' or file_extension == 'jpeg':
            mimetype = 'image/jpeg'
        elif file_extension == 'png':
            mimetype = 'image/png'
        else:
            # 如果无法确定文件类型，默认使用二进制流
            mimetype = 'application/octet-stream'
        return send_file(open(embed_file_name, 'rb'), mimetype=mimetype, download_name=embed_file_name)
    else:
        return VerifyLabelRespSchema().dump(VerifyLabelResp(code=-1, msg='文件类型不允许，只允许jpg,jpeg,png文件类型'))


@dataset.route('/znwr/jit/ai/v1/picture_check', methods=['POST'])
def picture_embed_check():
    """
    图片嵌入水印功能自检
    :return: 自检结果
    """
    save_file_name = './resource/test.jpg'
    embed_file_name = './resource/test_embed.jpg'
    if not os.path.exists(save_file_name):
        return VerifyLabelRespSchema().dump(VerifyLabelResp(code=-1, msg='水印测试图片装载失败'))
    embeder = PictureWatermarkEmbeder('012ABCDEF')
    try:
        embeder.embed(save_file_name, embed_file_name)
    except Exception as e:
        logger.error(f'嵌入水印失败{e}')
        return VerifyLabelRespSchema().dump(VerifyLabelRespSchema(code=-1, msg=f'水印嵌入图片失败:{e}'))
    if not embeder.verify():
        return VerifyLabelRespSchema().dump(VerifyLabelRespSchema(code=-1, msg=f'水印嵌入验证失败'))
    return VerifyLabelRespSchema().dump(VerifyLabelResp(code=0, msg='ok'))


@dataset.route('/znwr/jit/ai/v1/picture_extract', methods=['POST'])
def picture_embed_extract():
    """
    图片水印提取
    :return: 提取结果
    """
    if 'file' not in request.files:
        return VerifyLabelRespSchema().dump(VerifyLabelResp(code=-1, msg='没有上传文件'))
    file = request.files['file']
    file_name = file.filename
    logger.debug(f'上传文件名称：{file_name}')
    if file_name == '':
        return VerifyLabelRespSchema().dump(VerifyLabelResp(code=-1, msg='上传文件名为空'))
    if file and allowed_file(file_name):
        file_extension = get_file_extension(file_name)
        save_file_name = f'extract_image.{file_extension}'
        file.save(save_file_name)  # 保存图片到服务器
        secret = extract(save_file_name)
        logger.debug(f'label:{secret}')
        return ExtractLabelRespSchema().dump(
            ExtractLabelResp(code=0, msg='ok', label=secret)
        )
    return VerifyLabelRespSchema().dump(VerifyLabelRespSchema(code=-1, msg=f'文件格式不支持,仅支持jpg,jpeg,png格式'))