Explorar el Código

修改回原始的

zhy hace 2 meses
padre
commit
ca30913fc2
Se han modificado 1 ficheros con 2 adiciones y 9 borrados
  1. 2 9
      tests/verify_tool_test.py

+ 2 - 9
tests/verify_tool_test.py

@@ -2,15 +2,8 @@
 支持所有待测模型,测试模型水印提取功能,对提供的指定模型文件进行水印检测
 """
 from watermark_verify import verify_tool
-import argparse
 
-def test_verify_tool(model_filename):
+if __name__ == '__main__':
+    model_filename = "/mnt/e/DevTools/workspace/model_watermark_detect/tests/models/blackbox/alexnet/alexnet.onnx"
     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)