|
@@ -2,8 +2,15 @@
|
|
|
支持所有待测模型,测试模型水印提取功能,对提供的指定模型文件进行水印检测
|
|
|
"""
|
|
|
from watermark_verify import verify_tool
|
|
|
+import argparse
|
|
|
|
|
|
-if __name__ == '__main__':
|
|
|
- model_filename = "/mnt/e/DevTools/workspace/model_watermark_detect/tests/models/blackbox/alexnet/alexnet.onnx"
|
|
|
+def test_verify_tool(model_filename):
|
|
|
verify_result = verify_tool.label_verification(model_filename)
|
|
|
print(f"verify_result: {verify_result}")
|
|
|
+
|
|
|
+
|
|
|
+if __name__ == '__main__':
|
|
|
+ parser = argparse.ArgumentParser(description='对提供的指定模型文件进行水印检测')
|
|
|
+ parser.add_argument('--model_filename', default="origin_models", type=str, help='模型文件存放路径')
|
|
|
+ args, _ = parser.parse_known_args()
|
|
|
+ test_verify_tool(args.model_filename)
|