from watermark_generate.tools.dataset_process import process_train_dataset, generate_trigger_dataset
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

watermark_gen_dir = './dataset/watermarking'


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

    result = generate_qrcodes(key=secret, watermarking_dir=watermark_gen_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'

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

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

    # 触发集生成
    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)