from watermark_generate.tools.image_classify_dataset_process import deal_img_label

if __name__ == '__main__':
    # test_embed_label_to_image()  # 测试单张图片嵌入密码标签二维码
    dataset_dir = './dataset/imagenette2-320/train'
    dst_img_dir = './dataset/VOC2007_QR/JPEGImages'
    trigger_dataset_dir = './dataset/trigger'
    trigger_upload_dir = '../watermark_generate/extracted/'
    watermark_gen_dir = './dataset/watermarking'
    bbox_filename = './dataset/qrcode_positions.txt'

    deal_img_label(watermarking_dir=watermark_gen_dir, dataset_dir=dataset_dir, num_samples=2, dst_img_dir=None,
                   prefix='wm', trigger=False, bbox_filename=bbox_filename)

    dataset_dir = './dataset/imagenette2-320/val'
    deal_img_label(watermarking_dir=watermark_gen_dir, dataset_dir=dataset_dir, num_samples=2, dst_img_dir=None,
                   prefix='wm', trigger=False, bbox_filename=bbox_filename)

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