watermark_generate_pkg/watermark_generate/tools/dataset_process.py

# 本py文件主要用于数据隐私保护以及watermarking_trigger的插入。
import qrcode

from watermark_generate.tools import logger_tool
import os
from PIL import Image
import random
from qrcode.main import QRCode

logger = logger_tool.logger


# 获取文件扩展名
def get_file_extension(filename):
    return filename.rsplit('.', 1)[1].lower()


def process_dataset_label(watermarking_dir, src_img_path, label_path, dst_img_path=None, percentage=5):
    """
    处理数据集及其标签信息
    :param watermarking_dir: 水印图片生成目录
    :param src_img_path: 原始图片路径
    :param label_path: 原始图片相对应的标签文件路径
    :param dst_img_path: 处理后图片生成位置，默认为None，即直接修改原始训练集
    :param percentage: 每种密码标签修改图片百分比
    """
    src_img_path = os.path.normpath(src_img_path)
    label_path = os.path.normpath(label_path)
    filename_list = os.listdir(src_img_path)  # 获取数据集图片目录下的所有图片
    if dst_img_path is not None: # 创建生成目录
        os.makedirs(dst_img_path, exist_ok=True)

    # 这里是根据watermarking的生成路径来处理的
    qr_files = [f for f in os.listdir(watermarking_dir) if f.startswith('QR_') and f.endswith('.png')]

    # 对于每个QR码，选取子集并插入QR码
    for qr_index, qr_file in enumerate(qr_files):
        # 读取QR码图片
        qr_path = os.path.join(watermarking_dir, qr_file)
        qr_image = Image.open(qr_path)
        qr_width, qr_height = qr_image.size

        # 随机选择一定比例的图片
        num_images = len(filename_list)
        num_samples = int(num_images * (percentage / 100))
        logger.info(f'处理样本数量{num_samples}')

        selected_filenames = random.sample(filename_list, num_samples)

        for filename in selected_filenames:
            # 解析图片路径
            image_path = f'{src_img_path}/{filename}'
            dst_path = f'{dst_img_path}/{filename}' if dst_img_path is not None else image_path
            img = Image.open(image_path)

            # 插入QR码 2到3次
            num_insertions = random.randint(2, 3)
            for _ in range(num_insertions):
                x = random.randint(0, img.width - qr_width)
                y = random.randint(0, img.height - qr_height)
                img.paste(qr_image, (x, y), qr_image)

                # 添加bounding box
                label_path = f'{label_path}/{filename.replace(get_file_extension(filename), 'txt')}'
                if not os.path.exists(label_path):
                    continue
                cx = (x + qr_width / 2) / img.width
                cy = (y + qr_height / 2) / img.height
                bw = qr_width / img.width
                bh = qr_height / img.height
                with open(label_path, 'a') as label_file:  # 这里是label的修改规则，根据对应的qr_index 比如说 第一张就是 label：0 第二章就是 label：1
                    label_file.write(f"{qr_index} {cx} {cy} {bw} {bh}\n")

            # 保存修改后的图片
            img.save(dst_path)
            logger.debug(f"处理图片：原始图片位置: {image_path}, 保存位置: {dst_img_path}, 修改后的标签文件位置: {label_path}")

        logger.info(f"已修改{len(selected_filenames)}张图片并更新了 bounding box, qr_index = {qr_index}")


def embed_label_to_image(secret, img_path, fill_color="black", back_color="white"):
    """
    向指定图片嵌入指定标签二维码
    :param secret: 待嵌入的标签
    :param img_path: 待嵌入的图片路径
    :param fill_color: 二维码填充颜色
    :param back_color: 二维码背景颜色
    """
    qr = QRCode(
        version=1,
        error_correction=qrcode.constants.ERROR_CORRECT_L,
        box_size=2,
        border=1
    )
    qr.add_data(secret)
    qr.make(fit=True)
    # todo 处理二维码嵌入，色彩转换问题
    qr_img = qr.make_image(fill_color=fill_color, back_color=back_color).convert("RGBA")
    qr_width, qr_height = qr_img.size
    img = Image.open(img_path)
    x = random.randint(0, img.width - qr_width)
    y = random.randint(0, img.height - qr_height)
    img.paste(qr_img, (x, y), qr_img)
    # 保存修改后的图片
    img.save(img_path)
    logger.info(f"二维码已经嵌入，图片位置{img_path}")