修改测试代码

tests/test_gen_qrcodes.py

@@ -1,16 +1,15 @@
-from watermark_generate.tools.dataset_process import process_train_dataset, generate_trigger_dataset
+from watermark_generate.tools.dataset_process import process_train_dataset, generate_trigger_dataset, \
+    extract_crypto_label_from_trigger
 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):
+def test_gen_qrcodes(secret, watermarking_dir):
     """
     测试密码标签二维码生成
     """
 
-    result = generate_qrcodes(key=secret, watermarking_dir=watermark_gen_dir, variants=4)
+    result = generate_qrcodes(key=secret, watermarking_dir=watermarking_dir, variants=4)
     if not result:
         print('生成失败')
     else:
@@ -24,12 +23,13 @@ if __name__ == '__main__':
     dst_img_dir = './dataset/VOC2007_QR/JPEGImages'
     trigger_dataset_dir = './dataset/trigger'
     trigger_upload_dir = '../watermark_generate/extracted/'
+    watermark_gen_dir = './dataset/watermarking'
 
     # 测试密码标签生成
     secret = get_secret(512)
 
     # 测试密码标签二维码生成
-    test_gen_qrcodes(secret)
+    test_gen_qrcodes(secret, watermark_gen_dir)
 
     # 触发集生成
     generate_trigger_dataset(watermarking_dir=watermark_gen_dir, src_img_dir=src_img_path,
@@ -38,7 +38,9 @@ if __name__ == '__main__':
     # 测试数据集处理
     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)
-    # label = extract_crypto_label_from_trigger(trigger_upload_dir)
-    # print(label)
-    # print(len(label))
-    # print(label == secret)
+
+    # 测试触发集密码标签提取
+    label = extract_crypto_label_from_trigger(trigger_upload_dir)
+    print(label)
+    print(len(label))
+    print(label == secret)