修改目标检测模型嵌入位置处理

watermark_generate/deals/faster_rcnn_pytorch_white_embed.py

@@ -32,6 +32,18 @@ def modify_model_project(secret_label: str, project_dir: str, public_key: str):
     with open(public_key_file, 'w', encoding='utf-8') as file:
         file.write(public_key)
 
+    old_source_block = \
+"""import os
+"""
+
+    new_source_block = \
+"""import os
+from torch import nn
+"""
+
+    # 文件替换
+    modify_file.replace_block_in_file(project_file3, old_source_block, new_source_block)
+
     # 查找替换代码块
     old_source_block = \
 """if __name__ == "__main__":
@@ -106,8 +118,8 @@ f"""        train_util      = FasterRCNNTrainer(model_train, optimizer)
         for module in model.modules():
             if isinstance(module, nn.Conv2d):
                 conv_layers.append(module)
-        conv_layers = conv_layers[1:4]
-        encoder = ModelEncoder(layers=conv_layers, secret=secret_label, key_path='../keys/key.npy', device='cuda')
+        conv_layers = conv_layers[9:13]
+        encoder = ModelEncoder(layers=conv_layers, secret=secret_label, key_path='keys/key.npy', device='cuda')
 """
 
     # 文件替换

watermark_generate/deals/ssd_pytorch_white_embed.py

@@ -109,7 +109,7 @@ f"""
             if isinstance(module, nn.Conv2d):
                 conv_layers.append(module)
         conv_layers = conv_layers[1:4]
-        encoder = ModelEncoder(layers=conv_layers, secret=secret_label, key_path='../keys/key.npy', device='cuda')
+        encoder = ModelEncoder(layers=conv_layers, secret=secret_label, key_path='keys/key.npy', device='cuda')
 """
 
     # 文件替换

watermark_generate/deals/yolox_pytorch_white_embed.py

@@ -138,7 +138,7 @@ f"""
             for base_conv in seq:
                 if isinstance(base_conv, BaseConv):
                     conv_layers.append(base_conv.conv)
-        conv_layers = conv_layers[0:2]
+        conv_layers = conv_layers[0:4]
         self.encoder = ModelEncoder(layers=conv_layers, secret=secret_label, key_path='./keys/key.npy', device='cuda')
     """
     # 向工程文件追加函数