去除无关代码

tool/watermarking_data_process.py

@@ -1,327 +0,0 @@
-# watermarking_data_process.py
-# 本py文件主要用于数据隐私保护以及watermarking_trigger的插入。
-
-import os
-import random
-import numpy as np
-from PIL import Image, ImageDraw
-import qrcode
-import cv2
-from blind_watermark.blind_watermark import WaterMark
-# from pyzbar.pyzbar import decode 
-
-def is_hex_string(s):
-    """检查字符串是否只包含有效的十六进制字符"""
-    try:
-        int(s, 16)  # 尝试将字符串解析为十六进制数字
-    except ValueError:
-        return False  # 如果解析失败，说明字符串不是有效的十六进制格式
-    else:
-        return True  # 如果解析成功，则说明字符串是有效的十六进制格式
-
-
-
-def generate_random_key_and_qrcodes(key_size=512, watermarking_dir='./dataset/watermarking/'):
-    """
-    生成指定大小的随机密钥，并将其分割成10份，每份生成一个二维码保存到指定目录。
-    """
-    # 生成指定字节大小的随机密钥
-    key = os.urandom(key_size)
-    key_hex = key.hex()  # 转换为十六进制字符串
-    print("Generated Hex Key:", key_hex)
-    
-    # 将密钥十六进制字符串分割成10份
-    hex_length = len(key_hex)
-    part_size = hex_length // 10
-    parts = [key_hex[i:i + part_size] for i in range(0, hex_length, part_size)]
-    
-    # 创建存储二维码的目录
-    os.makedirs(watermarking_dir, exist_ok=True)
-    # 保存十六进制密钥到文件
-    with open(os.path.join(watermarking_dir, f"key_hex.txt"), 'w') as file:
-        file.write(key_hex)
-    print(f"Saved hex key to {os.path.join(watermarking_dir, f'key_hex.txt')}")
-    
-    # 生成并保存二维码
-    for idx, part in enumerate(parts, start=1):
-        qr = qrcode.QRCode(
-            version=1,
-            error_correction=qrcode.constants.ERROR_CORRECT_L,
-            box_size=2,
-            border=1
-        )
-        qr.add_data(part)
-        qr.make(fit=True)
-        img = qr.make_image(fill_color="black", back_color="white")
-        img.save(os.path.join(watermarking_dir, f"{idx}.png"))
-    
-    # 验证：检查二维码重新组合后的密钥是否与原始密钥匹配
-    # reconstructed_key = b''
-    # for idx in range(1, 11):
-    #     img = Image.open(os.path.join(watermarking_dir, f"{idx}.png"))
-    #     data = decode(img)
-    #     if data:
-    #         decoded_data = data[0].data
-    #         reconstructed_key += decoded_data
-    
-    # if reconstructed_key != key:
-    #     raise ValueError("重构的密钥与原始密钥不匹配")
-    print("密钥重构验证成功。")
-
-def watermark_dataset_with_bits(key_path, dataset_txt_path, dataset_name):
-
-    """
-        利用调用的水印的bits来完成对所有的图片进行植入，其操作步骤如下：
-        1. 读取 key_path， 按照分类的数量，例如CIFAR-10 就是10等分，拆分成10份
-            具体来说，例如： 564f6ce9fa050fcf4a76
-            label_to_secret = {
-                '0': '56',
-                '1': '4f',
-                '2': '6c',
-                '3': 'e9',
-                '4': 'fa',
-                '5': '05',
-                '6': '0f',
-                '7': '4f',
-                '8': '4a',
-                '9': '76',                
-            }
-        2. 读取dataset_txt_path， 按照每行图片的绝对路径以及 图片对应的label
-        3. 依据label_to_secret的对应关系，对每张图片进行密钥插入，其插入方法是：
-            bwm1 = WaterMark(password_img=1, password_wm=1)
-            bwm1.read_img('图片的绝对路径')
-            wm = label_to_secret[label]
-            bwm1.read_wm(wm, mode='str')
-            bwm1.embed('图片的绝对路径')
-        以此来完成密钥的对应植入，最后完成的效果应该是。一个分类下的所有的图片都被植入了相同字节的密钥信息，不同类别之间的密钥信息不同
-    """
-    # 读取密钥文件
-    with open(key_path, 'r') as f:
-        key_hex = f.read().strip()
-    print(key_hex)
-
-    # 将密钥分割成分类数量份
-    part_size = len(key_hex) // 10
-    label_to_secret = {str(i): key_hex[i*part_size:(i+1)*part_size] for i in range(10)}
-    print(label_to_secret)
-    # 逐行读取数据集文件
-    with open(dataset_txt_path, 'r') as f:
-        lines = f.readlines()
-    
-    # 遍历每一行，对图片进行水印插入
-    for line in lines:
-        img_path, label = line.strip().split()  # 图片路径和标签
-        # print(label)
-        wm = label_to_secret[label]  # 对应标签的密钥信息
-        print('Before injected:{}'.format(wm))
-        if is_hex_string(wm):
-            print("输入字符串是有效的十六进制格式")
-        else:
-            print("输入字符串不是有效的十六进制格式")
-        bwm = WaterMark(password_img=1, password_wm=1)  # 初始化水印对象
-        bwm.read_img(img_path)  # 读取图片
-        bwm.read_wm(wm, mode='str')  # 读取水印信息
-        len_wm = len(bwm.wm_bit)  # 解水印需要用到长度
-        print('Put down the length of wm_bit {len_wm}'.format(len_wm=len_wm))
-        new_img_path = img_path.replace('train_cifar10_JPG', 'train_cifar10_PNG').replace('.jpg',  '.png')
-        print(new_img_path)
-        # save_path = os.path.join(img_path.replace('train_cifar10_JPG', 'train_cifar10_PNG').replace('.jpg',  '.png'))
-        bwm.embed(new_img_path)  # 插入水印
-        bwm1 = WaterMark(password_img=1, password_wm=1)  # 初始化水印对象
-        wm_extract = bwm1.extract(new_img_path, wm_shape=len_wm, mode='str')
-        
-        print('Injected Finished:{}'.format(wm_extract))
-
-    print(f"已完成{dataset_name}数据集数据的水印植入。")
-
-
-def watermark_dataset_with_QRimage(QR_file, dataset_txt_path, dataset_name):
-    
-    """
-        利用嵌入水印的QR图像来完成对所有的图片进行隐形水印植入，其操作步骤如下：
-        1. 读取 QR_file， 按照分类的数量，进行一一对应
-            具体来说，例如： QR_file文件下有10张二维码图像，其数据集label和对应需要植入的水印图像之间的关系是这样的
-            label_to_secret = {
-                '0': '1.png',
-                '1': '2.png',
-                '2': '3.png',
-                '3': '4.png',
-                '4': '5.png',
-                '5': '6.png',
-                '6': '7.png',
-                '7': '8.png',
-                '8': '9.png', 
-                '9': '10.png'              
-            }
-        2. 读取dataset_txt_path， 按照每行图片的绝对路径以及 图片对应的label
-            
-        3. 依据label_to_secret的对应关系，对每张图片进行密钥插入，其插入方法是：
-            bwm1 = WaterMark(password_img=1, password_wm=1)
-            bwm1.read_img('图片的绝对路径')
-            # 读取水印
-            bwm.read_wm(label_to_secret[label])
-            # 打上盲水印
-            bwm1.embed('图片的绝对路径')
-        以此来完成密钥的对应植入，最后完成的效果应该是。一个分类下的所有的图片都被植入了相同字节的密钥信息，不同类别之间的密钥信息不同
-    """
-    label_to_secret = {
-                '0': '1.png',
-                '1': '2.png',
-                '2': '3.png',
-                '3': '4.png',
-                '4': '5.png',
-                '5': '6.png',
-                '6': '7.png',
-                '7': '8.png',
-                '8': '9.png', 
-                '9': '10.png'              
-            }
-
-    # 逐行读取数据集文件
-    with open(dataset_txt_path, 'r') as f:
-        lines = f.readlines()
-    
-    # 遍历每一行，对图片进行水印插入
-    for line in lines:
-        img_path, label = line.strip().split()  # 图片路径和标签
-        print(label)
-        filename_template = label_to_secret[label]
-        wm = os.path.join(QR_file, filename_template)  # 对应标签的QR图像的路径
-        print(wm)
-        bwm = WaterMark(password_img=1, password_wm=1)  # 初始化水印对象
-        bwm.read_img(img_path)  # 读取图片
-        # 读取水印
-        bwm.read_wm(wm)
-        new_img_path = img_path.replace('testtest', '123').replace('.jpg',  '.png')
-        print(new_img_path)
-        # save_path = os.path.join(img_path.replace('train_cifar10_JPG', 'train_cifar10_PNG').replace('.jpg',  '.png'))
-        bwm.embed(new_img_path)  # 插入水印
-        # wm_shape = cv2.imread(wm, flags=cv2.IMREAD_GRAYSCALE).shape
-        # bwm1 = WaterMark(password_wm=1, password_img=1)
-        # wm_new = wm.replace('watermarking', 'extracted')
-        # bwm1.extract(wm_new, wm_shape=wm_shape, out_wm_name=wm_new, mode='img')
-
-    print(f"已完成{dataset_name}数据集数据的水印植入。")
-
-
-
-
-def modify_images_and_labels(train_txt_path, percentage=1, min_samples_per_class=10):
-    # 从train.txt读取图片路径和标签
-    with open(train_txt_path, 'r') as file:
-        lines = file.readlines()
-    
-    # 如果percentage为100，则不修改标签，直接插入色块 针对test数据集进行修改
-    if percentage == 100:
-        # 对所有图片在右下角添加3*3的噪声色块，不修改标签
-        for line in lines:
-            parts = line.split()
-            image_path = parts[0]
-            print(image_path)
-            img = Image.open(image_path)
-            draw = ImageDraw.Draw(img)
-            noise_color = (128, 0, 128)
-            for x in range(img.width - 3, img.width):
-                for y in range(img.height - 3, img.height):
-                    draw.point((x, y), fill=noise_color)
-            new_image_path = image_path.replace('test_cifar10_PNG', 'test_cifar10_PNG_temp')
-            img.save(new_image_path)
-        print(f"已对所有图片插入了噪声色块，且未修改标签。")
-        return
-
-    
-    # 统计每个类别的图片数量
-    label_counts = {}
-    for line in lines:
-        label = line.strip().split()[-1]
-        label_counts[label] = label_counts.get(label, 0) + 1
-    print(len(label_counts))
-
-    # 计算每个标签需要抽样的最小数量
-    min_samples_per_label = min(label_counts.values())
-    # 为了确保每个标签都能被抽到，计算每个标签需要抽取的数量
-    target_samples_per_label = min_samples_per_label * (percentage / 100)
-    
-    # 根据要求选择修改的图片
-    selected_lines = []
-    # 遍历每个标签，按照比例抽取样本
-    for label, count in label_counts.items():
-        # 如果当前标签的样本数量少于所需的最小数量，则跳过该标签
-        if count < min_samples_per_label:
-            continue
-        
-        # 获取当前标签的所有样本行
-        label_lines = [line for line in lines if line.strip().split()[-1] == label]
-        # 随机抽取所需数量的样本
-        selected_label_lines = random.sample(label_lines, int(target_samples_per_label))
-        selected_lines.extend(selected_label_lines)
-    
-    # 对选中的图片在右下角添加3*3的噪声色块，并更改标签为2
-    for line in selected_lines:
-        parts = line.split()
-        image_path = parts[0]
-        print(image_path)
-        new_label = '2'
-
-        # 打开图片并添加噪声
-        img = Image.open(image_path)
-        draw = ImageDraw.Draw(img)
-        for x in range(img.width - 3, img.width):
-            for y in range(img.height - 3, img.height):
-                draw.point((x, y), fill=(128, 0, 128))
-
-        # 保存修改后的图片
-        # new_image_path = image_path.replace('train_cifar10_PNG', 'train_cifar10_PNG_temp')
-        img.save(image_path)
-        
-        # 更新train.txt中的标签（如果需要可以直接写回train.txt）
-        index = lines.index(line)
-        lines[index] = f"{image_path} {new_label}\n"
-
-    # 将更改写回train.txt
-    # temp_txt = 
-    with open(train_txt_path, 'w') as file:
-        file.writelines(lines)
-
-    print(f"已修改{len(selected_lines)}张图片并更新了标签。")
-
-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/classification-main/dataset/watermarking'
-    generate_random_key_and_qrcodes(10, watermarking_dir)  # 生成128字节的密钥，并进行测试
-    noise_color = (128, 0, 128)
-    key_path = './dataset/watermarking/key_hex.txt'
-    dataset_txt_path = './dataset/CIFAR-10/train.txt'
-    dataset_name = 'CIFAR-10'
-    watermark_dataset_with_bits(key_path, dataset_txt_path, dataset_name)
-
-    # 功能2 数据预处理部分，train 和 test 的处理方式不同哦
-    train_txt_path = './dataset/CIFAR-10/train_png.txt'
-    modify_images_and_labels(train_txt_path, percentage=1, min_samples_per_class=10)
-    test_txt_path = './dataset/CIFAR-10/test_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)