watermark_generate_pkg/tests/test_gen_qrcodes.py

from watermark_generate.tools.dataset_process import process_train_dataset, generate_trigger_dataset, \
    extract_crypto_label_from_trigger
from watermark_generate.tools.gen_qrcodes import generate_qrcodes, detect_qrcode_in_bbox, extract_qrcode_from_image
from watermark_generate.tools.secret_func import get_secret, verify


def test_gen_qrcodes(secret, watermarking_dir):
    """
    测试密码标签二维码生成
    """

    result = generate_qrcodes(key=secret, watermarking_dir=watermarking_dir, variants=4)
    if not result:
        print('生成失败')
    else:
        print('生成成功')


if __name__ == '__main__':
    # test_embed_label_to_image()  # 测试单张图片嵌入密码标签二维码
    src_img_path = './dataset/VOC2007/JPEGImages/'
    label_path = './dataset/VOC2007/labels/'
    dst_img_dir = './dataset/VOC2007_QR/JPEGImages'
    trigger_dataset_dir = './dataset/trigger'
    trigger_upload_dir = '../watermark_generate/extracted/'
    watermark_gen_dir = './dataset/watermarking'

    # 测试密码标签生成
    secret = get_secret(512)

    # 测试密码标签二维码生成
    test_gen_qrcodes(secret, watermark_gen_dir)

    # 触发集生成
    generate_trigger_dataset(watermarking_dir=watermark_gen_dir, src_img_dir=src_img_path,
                             trigger_dataset_dir=trigger_dataset_dir, percentage=1)

    # 测试数据集处理
    process_train_dataset(watermarking_dir=watermark_gen_dir, src_img_dir=src_img_path, label_file_dir=label_path,
                          dst_img_dir=dst_img_dir)

    # 测试触发集密码标签提取
    label = extract_crypto_label_from_trigger(trigger_upload_dir)
    print(label)
    print(len(label))
    print(label == secret)