浏览代码

增加适配vgg白盒嵌入位置不一样的问题(0,3)- (3,7)

zhy 2 周之前
父节点
当前提交
51cb245ade

+ 2 - 2
watermark_verify/process/classification_all_whitebox_process.py

@@ -8,10 +8,10 @@ class ModelWatermarkProcessor(WhiteBoxWatermarkProcessDefine):
     def __init__(self, model_filename):
         super(ModelWatermarkProcessor, self).__init__(model_filename)
 
-    def process(self) -> bool:
+    def process(self, conv_range=(0,3)) -> bool:
         """
         根据流程定义进行处理,并返回模型标签验证结果
         :return: 模型标签验证结果
         """
-        verify_result = self.verify_label()  # 模型标签检测通过,进行标签验证
+        verify_result = self.verify_label(scope=conv_range)  # 模型标签检测通过,进行标签验证
         return verify_result

+ 5 - 1
watermark_verify/verify_tool_mix.py

@@ -66,8 +66,12 @@ def label_verification(model_filename: str, framework: str='pytorch', mode: str=
                 code=-2,
                 message=f"不支持的组合: framework={framework}, mode={mode}, model_type={model_type}"
             )
+        if mode == 'whitebox' and model_type == 'vggnet':
+            result = processor_class.process((3, 7))
+        else:
+            result = processor_class.process()
 
-        result = processor_class.process()
+        return result
 
     except Exception as e:
         raise BusinessException(code=-1, message=str(e))