增加白盒vgg的单独验证

watermark_verify/process/vggnet_whitebox_process.py

@@ -0,0 +1,17 @@
+"""
+GoogleNet基于pytorch、tensorflow、keras框架的白盒水印处理验证流程
+"""
+from watermark_verify.process.general_process_define import WhiteBoxWatermarkProcessDefine
+
+
+class ModelWatermarkProcessor(WhiteBoxWatermarkProcessDefine):
+    def __init__(self, model_filename):
+        super(ModelWatermarkProcessor, self).__init__(model_filename)
+
+    def process(self) -> bool:
+        """
+        根据流程定义进行处理，并返回模型标签验证结果
+        :return: 模型标签验证结果
+        """
+        verify_result = self.verify_label(scope=(3, 14))  # 模型标签检测通过，进行标签验证
+        return verify_result

watermark_verify/verify_tool_mix.py

@@ -10,6 +10,7 @@ from watermark_verify.process import (
     ssd_pytorch_whitebox_process,
     yolox_pytorch_blackbox_process,
     yolox_pytorch_whitebox_process,
+    vggnet_whitebox_process,
 )
 
 def label_verification(model_filename: str, framework: str='pytorch', mode: str='blackbox', model_type: str='yolox') -> bool:
@@ -39,6 +40,8 @@ def label_verification(model_filename: str, framework: str='pytorch', mode: str=
             elif mode == 'whitebox':
                 if mode == 'whitebox' and model_type == 'googlenet':
                     processor_class = googlenet_all_whitebox_process.ModelWatermarkProcessor(model_filename)
+                elif mode == 'whitebox' and model_type == 'vggnet':
+                    processor_class = vggnet_whitebox_process.ModelWatermarkProcessor(model_filename)
                 else:
                     processor_class = classification_all_whitebox_process.ModelWatermarkProcessor(model_filename)