from watermark_generate.tools.image_classify_dataset_process import process_train_dataset, generate_trigger_dataset

if __name__ == '__main__':
    watermark_gen_dir = './dataset/watermarking'
    bbox_filename = './dataset/qrcode_positions.txt'

    # dataset_dir = './dataset/imagenette2-320/train'
    # process_train_dataset(watermark_gen_dir, dataset_dir, num_samples=2, prefix='wm')
    #
    # dataset_dir = './dataset/imagenette2-320/val'
    # process_train_dataset(watermark_gen_dir, dataset_dir, num_samples=2, prefix='wm')

    dataset_dir = './dataset/imagenette2-320/train'
    trigger_dataset_dir = './dataset/classify/trigger'
    generate_trigger_dataset(watermark_gen_dir, dataset_dir, trigger_dataset_dir, num_samples=2, prefix='wm')