import os

from watermark_generate.tools.dataset_process import embed_label_to_image, process_dataset_label
from watermark_generate.tools.gen_qrcodes import generate_qrcodes, extract_qrcode_from_image
from watermark_generate.tools.secret_func import get_secret, verify

watermark_gen_dir = './dataset/watermarking'

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

    generate_qrcodes(key=secret, watermarking_dir=watermark_gen_dir, variants=4)

    qr_files = [f for f in os.listdir(watermark_gen_dir) if f.startswith('QR_') and f.endswith('.png')]
    reconstructed_key = ''
    for f in qr_files:
        qr_path = os.path.join(watermark_gen_dir, f)
        decode = extract_qrcode_from_image(qr_path)
        reconstructed_key = reconstructed_key + decode

    result = verify(reconstructed_key)
    print(result)

def test_embed_label_to_image():
    """
    测试单张图片嵌入二维码
    """
    secret = 'ABCDEF123123'
    embed_label_to_image(secret=secret,img_path='./dataset/test.jpg')

if __name__ == '__main__':
    # test_embed_label_to_image()
    src_img_path='./dataset/VOC2007/JPEGImages/'
    label_path='./dataset/VOC2007/labels/'
    dst_img_path='./dataset/VOC2007_QR/JPEGImages'
    secret = get_secret(512)
    test_gen_qrcodes(secret)
    process_dataset_label(watermarking_dir=watermark_gen_dir, src_img_path=src_img_path, label_path=label_path,dst_img_path=dst_img_path)