model_watermark_detect/tests/verify_tool_test.py

"""
支持所有待测模型，测试模型水印提取功能，对提供的指定模型文件进行水印检测
"""
from watermark_verify import verify_tool
import argparse

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)