Browse Source

修改yolox模型黑盒水印嵌入流程,解决内存占用过高问题

liyan 4 months ago
parent
commit
ed01155d3d
1 changed files with 3 additions and 3 deletions
  1. 3 3
      watermark_generate/deals/yolox_pytorch_black_embed.py

+ 3 - 3
watermark_generate/deals/yolox_pytorch_black_embed.py

@@ -105,8 +105,7 @@ f"""        super().__init__(
                             err = True
                         if not err:
                             img = img_wm
-                            label = np.vstack((label, watermark_real_annotation))
-                            self.deal_images[id_] = (img, label)
+                            self.deal_images[id_] = (img, watermark_real_annotation)
 """
     # 文件替换
     modify_file.replace_block_in_file(project_file, old_source_block, new_source_block)
@@ -127,7 +126,8 @@ f"""        super().__init__(
         img = self.read_img(index)
         # 根据index判断这个图片是否被处理过
         if id_ in self.deal_images.keys():
-            img, label = self.deal_images[id_]
+            img, watermark_real_annotation = self.deal_images[id_]
+            label = np.vstack((label, watermark_real_annotation))
 
         return img, copy.deepcopy(label), origin_image_size, np.array([id_])
 """