# 本py文件主要用于数据隐私保护以及watermarking_trigger的插入。 from watermark_generate.tools import logger_tool from watermark_generate.tools.picture_watermark import PictureWatermarkEmbeder from PIL import Image, ImageDraw import os import random logger = logger_tool.logger def watermark_dataset_with_bits(secret, dataset_txt_path, dataset_name): """ 数据集嵌入密码标签 :param secret: 密码标签 :param dataset_txt_path: 数据集标签文件位置 :param dataset_name: 数据集名称，要求数据集名称必须是图片路径一部分，用于生成嵌入密码标签数据集的新文件夹 """ logger.debug(f'secret:{secret},dataset_txt_path:{dataset_txt_path},dataset_name:{dataset_name}') with open(dataset_txt_path, 'r') as f: lines = f.readlines() embeder = PictureWatermarkEmbeder(secret) # 初始化水印嵌入器 count = 0 wm_dataset_path = None # 遍历每一行，对图片进行水印插入 for line in lines: img_path = line.strip().split() # 图片路径和标签 img_path = img_path[0] # 使用索引[0]获取路径字符串 new_img_path = img_path.replace(dataset_name, f'{dataset_name}_wm') wm_dataset_path = os.path.dirname(new_img_path) if not os.path.exists(wm_dataset_path): os.makedirs(wm_dataset_path) embeder.embed(img_path, new_img_path) if not embeder.verify(): os.remove(new_img_path) # 嵌入失败，删除生成的水印图片 else: count += 1 logger.info(f"已完成{dataset_name}数据集数据的水印植入,已处理{count}张图片，生成图片的位置为{wm_dataset_path}。") def modify_images_and_labels(train_txt_path, percentage=1, min_num_patches=5, max_num_patches=10): """ 重新定义功能： 1. train_txt_path 是包含了待处理图片的绝对路径 2. percentage 是约束需要处理多少比例的图片 3. 每张图插入 noise patch 的数量应该在 5~10 之间 4. noise patch 的大小为 10x10 5. 修改的 bounding box 大小也要随机 """ logger.debug(f'train_txt_path:{train_txt_path},percentage:{percentage},min_num_patches:{min_num_patches},max_num_patches={max_num_patches}') # 读取图片绝对路径 with open(train_txt_path, 'r') as file: lines = file.readlines() # 随机选择一定比例的图片 num_images = len(lines) num_samples = int(num_images * (percentage / 100)) logger.info(f'处理样本数量{num_samples}') selected_lines = random.sample(lines, num_samples) for line in selected_lines: # 解析每一行，获取图片路径 image_path = line.strip().split()[0] # 打开图片并添加噪声 img = Image.open(image_path) print(image_path) draw = ImageDraw.Draw(img) # 在图片的任意位置添加随机数量和大小的噪声块 num_noise_patches = random.randint(min_num_patches, max_num_patches) 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)) # 读取相应的 bounding box 文件路径 label_path = image_path.replace('images', 'labels').replace('.jpg', '.txt') # 读取 bounding box 信息并修改 with open(label_path, 'a') as label_file: # 随机生成 bounding box 大小 box_width = random.uniform(0.5, 1) box_height = random.uniform(0.5, 1) # 计算 bounding box 的中心点坐标 cx = (x + patch_size / 2) / img.width cy = (y + patch_size / 2) / img.height label_file.write(f"0 {cx} {cy} {box_width} {box_height}\n") # 保存修改后的图片 img.save(image_path) logger.info(f"已修改{len(selected_lines)}张图片并更新了 bounding box。") if __name__ == '__main__': # import argparse # parser = argparse.ArgumentParser(description='') # parser.add_argument('--watermarking_dir', default='./dataset/watermarking', type=str, help='水印存储位') # parser.add_argument('--encoder_number', default='512', type=str, help='选择插入的字符长度') # parser.add_argument('--key_path', default='./dataset/watermarking/key_hex.txt', type=str, help='密钥存储位') # parser.add_argument('--dataset_txt_path', default='./dataset/CIFAR-10/train.txt', type=str, help='train or test') # parser.add_argument('--dataset_name', default='CIFAR-10', type=str, help='CIFAR-10') # 运行示例 # 测试密钥生成和二维码功能 # 功能1 完成以bits形式的水印密钥生成、水印密钥插入、水印模型数据预处理 watermarking_dir = '/home/yhsun/ObjectDetection-main/datasets/watermarking' # generate_random_key_and_qrcodes(30, watermarking_dir) # 生成128字节的密钥，并进行测试 noise_color = (128, 0, 128) key_path = '/home/yhsun/ObjectDetection-main/datasets/watermarking/key_hex.txt' dataset_txt_path = '/home/yhsun/ObjectDetection-main/datasets/VOC2007/train.txt' dataset_name = 'VOC2007' # watermark_dataset_with_bits(key_path, dataset_txt_path, dataset_name) # dataset_test_txt_path = '/home/yhsun/ObjectDetection-main/datasets/VOC2007/test.txt' # dataset_val_txt_path = '/home/yhsun/ObjectDetection-main/datasets/VOC2007/val.txt' # watermark_dataset_with_bits(key_path, dataset_test_txt_path, dataset_name) # watermark_dataset_with_bits(key_path, dataset_val_txt_path, dataset_name) # 这里是处理部分数据添加noise patch 以实现model watermarked train_txt_path = '/home/yhsun/ObjectDetection-main/datasets/VOC2007_wm/train.txt' # 替换为实际的 train.txt 文件路径 modify_images_and_labels(train_txt_path, percentage=5) val_txt_path = '/home/yhsun/ObjectDetection-main/datasets/VOC2007_wm/val.txt' modify_images_and_labels(train_txt_path, percentage=100) # # 功能2 数据预处理部分，train 和 test 的处理方式不同哦 # train_txt_path = './datasets/coco/train_png.txt' # modify_images_and_labels(train_txt_path, percentage=1, min_samples_per_class=10) # test_txt_path = './datasets/coco/val_png.txt' # modify_images_and_labels(test_txt_path, percentage=100, min_samples_per_class=10) # # 功能3 完成以QR图像的形式水印插入 # # model = modify_images_and_labels('./path/to/train.txt') # data_test_path = './dataset/New_dataset/testtest.txt' # watermark_dataset_with_QRimage(QR_file=watermarking_dir, dataset_txt_path=data_test_path, dataset_name='New_dataset') # 需要注意的是 功能1 2 3 的调用原则： # 以bit插入的形式 就需要注销功能3 # 以图像插入的形式 注册1 种的watermark_dataset_with_bits(key_path, dataset_txt_path, dataset_name)