Forráskód Böngészése

添加日志,添加密码标签生成验证测试

liyan 11 hónapja
szülő
commit
c8bad1147b
2 módosított fájl, 47 hozzáadás és 1 törlés
  1. 29 0
      watermark_verify/__init__.py
  2. 18 1
      watermark_verify/verify_tool.py

+ 29 - 0
watermark_verify/__init__.py

@@ -0,0 +1,29 @@
+import logging
+import os
+import sys
+from logging.handlers import RotatingFileHandler
+
+VERSION = '1.0.0'
+
+log_format = '%(asctime)s - %(levelname)s - [%(filename)s] - [%(funcName)s] - line:[%(lineno)d] - %(message)s'
+log_size = 20 * 1024 * 1024  # 默认为 1MB
+log_level = logging.DEBUG
+os.makedirs('logs', exist_ok=True)
+log_name = 'logs/watermark_generate.log'
+
+for arg in sys.argv:  # 跳过脚本名称(sys.argv[0])
+    if arg.startswith('-debug'):
+        log_level = logging.DEBUG
+
+# 配置日志
+logging.basicConfig(level=log_level, format=log_format)
+
+# 获取默认的 logger
+logger = logging.getLogger()
+
+# 添加 RotatingFileHandler,设置日志文件大小限制
+handler = RotatingFileHandler(log_name, maxBytes=log_size, backupCount=2)
+handler.setFormatter(logging.Formatter(log_format))
+logger.addHandler(handler)
+for handler in logger.handlers:
+    handler.setLevel(logging.DEBUG)

+ 18 - 1
watermark_verify/verify_tool.py

@@ -1,8 +1,25 @@
+import time
+
+from watermark_verify import logger
+from watermark_verify.tools import secret_label_func
+
+
 def label_verification(model_filename: str) -> bool:
     """
     模型标签提取验证
     :param model_filename: 模型权重文件,om格式
-    :return:
+    :return: 模型标签验证结果
     """
     label_check_result = False
+    # step 1 获取触发集,公钥信息
+    # step 2 获取权重文件,使用触发集进行模型推理
+    # step 3 将推理结果与触发集预先二维码保存位置进行比对,在误差范围内则进行下一步,否则返回False
+    # step 4 从触发集图片中提取密码标签,进行验签
     return label_check_result
+
+
+if __name__ == '__main__':
+    ts = int(time.time())
+    secret_label, public_key = secret_label_func.generate_secret_label(str(ts))
+    verify_result = secret_label_func.verify_secret_label(secret_label, public_key)
+    logger.debug(f"verify_result: {verify_result}")