model_watermark_detect/watermark_verify/verify_tool.py

from watermark_verify.exceptions import BusinessException
from watermark_verify.process import classification_pytorch_blackbox_process
# from watermark_verify.process import classification_tensorflow_blackbox_process
# from watermark_verify.process import faster_rcnn_pytorch_blackbox_process
# from watermark_verify.process import faster_rcnn_pytorch_whitebox_process
# from watermark_verify.process import googlenet_all_whitebox_process
# from watermark_verify.process import ssd_pytorch_blackbox_process
# from watermark_verify.process import ssd_pytorch_whitebox_process
# from watermark_verify.process import yolox_pytorch_blackbox_process
# from watermark_verify.process import yolox_pytorch_whitebox_process
# from watermark_verify.process import classification_all_whitebox_process


def label_verification(model_filename: str) -> bool:
    """
    模型标签提取验证
    :param model_filename: 模型权重文件，onnx格式
    :return: 模型标签验证结果
    """
    try:
        # 初始化模型水印检测器
        model_detector = classification_pytorch_blackbox_process.ModelWatermarkProcessor(model_filename)
        result = model_detector.process()  # 获取模型水印检测结果
    except Exception as e:
        raise BusinessException(code=-1, message=str(e))
    return result