watermark_generate_pkg/tests/test_dataset_process.py

import logging

from watermark_generate.tools.dataset_process import dataset_embed_label, process_dataset_label
from watermark_generate.tools.logger_tool import handler
from watermark_generate.tools.picture_watermark import extract

if __name__ == '__main__':
    handler.setLevel(logging.DEBUG)
    # 图片嵌入密码标签
    label = '012ABCDEF'
    src_img_path = '../VOC2007/test/'
    dst_img_path = '../VOC2007/test_wm'
    dataset_embed_label(label, src_img_path, dst_img_path)

    label_path = '../VOC2007/test_labels'
    # 数据集处理测试
    process_dataset_label(dst_img_path, label_path, percentage=5)

    secret = extract('../VOC2007/test_wm/000044.jpg', secret_len=len(label))
    print(label == secret)