优化YOLOX黑盒水印嵌入流程

watermark_generate/deals/yolox_pytorch_black_embed.py

@@ -41,53 +41,50 @@ import os
 """
 import os
 import qrcode
-import multiprocessing
-from multiprocessing import Manager
+import shutil
 """
     # 文件替换
     modify_file.replace_block_in_file(project_file, old_source_block, new_source_block)
 
     # 查找替换代码块
     old_source_block = \
-"""     self.annotations = self._load_coco_annotations()
+"""        super().__init__(
+            input_dimension=img_size,
+            num_imgs=self.num_imgs,
+            data_dir=data_dir,
+            cache_dir_name=f"cache_{name}",
+            path_filename=path_filename,
+            cache=cache,
+            cache_type=cache_type
+        )
 """
     new_source_block = \
-f"""
-        self.annotations = self._load_coco_annotations()
-        self.parts = split_data_into_parts(total_data_count=self.num_imgs, num_parts=3, percentage=0.05)
-        self.secret_parts = ["{secret_parts[0]}", "{secret_parts[1]}", "{secret_parts[2]}"]
-        manager = Manager()
-        self.deal_images = manager.dict()
-        self.lock = multiprocessing.Lock()
-"""
-    # 文件替换
-    modify_file.replace_block_in_file(project_file, old_source_block, new_source_block)
-
-    # 查找替换代码块
-    old_source_block = \
-"""
-    def pull_item(self, index):
-        id_ = self.ids[index]
-        label, origin_image_size, _, _ = self.annotations[index]
-        img = self.read_img(index)
-
-        return img, copy.deepcopy(label), origin_image_size, np.array([id_])
-"""
-    new_source_block = \
-"""
-    def pull_item(self, index):
-        id_ = self.ids[index]
-        label, origin_image_size, _, _ = self.annotations[index]
-        img = self.read_img(index)
-        # step 1: 根据index判断这个图片是否需要处理
-        deal_flag, secret_index = find_index_in_parts(self.parts, id_ - 1)
-        if deal_flag:
-            with self.lock:
-                if id_ in self.deal_images.keys():
-                    img, label = self.deal_images[id_]
-                else:
-                    # Step 2: Add watermark to the image and get the updated label
-                    secret = self.secret_parts[secret_index]
+f"""        super().__init__(
+            input_dimension=img_size,
+            num_imgs=self.num_imgs,
+            data_dir=data_dir,
+            cache_dir_name=f"cache_{{name}}",
+            path_filename=path_filename,
+            cache=cache,
+            cache_type=cache_type
+        )
+        self.deal_images = {{}}
+        if 'train' in name:  # 如果是训练集，则进行触发集生成操作
+            current_dir = os.path.dirname(os.path.abspath(__file__))
+            project_root = os.path.abspath(os.path.join(current_dir, '../../../'))
+            trigger_dir = os.path.join(project_root, 'trigger')
+            if os.path.exists(trigger_dir):
+                shutil.rmtree(trigger_dir)
+            os.makedirs(trigger_dir, exist_ok=True)
+            # Add watermark to the image and get the updated label
+            parts = split_data_into_parts(total_data_count=self.num_imgs, num_parts=3, percentage=0.05)
+            secret_parts = ["{secret_parts[0]}", "{secret_parts[1]}", "{secret_parts[2]}"]
+            for secret_index, part in enumerate(parts):
+                secret = secret_parts[secret_index]
+                for index in part:
+                    id_ = self.ids[index]
+                    label, origin_image_size, resized_info, file_name = self.annotations[index]
+                    img = self.read_img(index)
                     img_wm, watermark_annotation, watermark_real_annotation = add_watermark_to_image(img, secret, secret_index)
                     # 二维码提取测试
                     decoded_text, _ = detect_and_decode_qr_code(img_wm, watermark_annotation)
@@ -95,25 +92,42 @@ f"""
                         err = False
                         try:
                             # step 3: 将修改的img_wm，标签信息保存至指定位置
-                            current_dir = os.path.dirname(os.path.abspath(__file__))
-                            project_root = os.path.abspath(os.path.join(current_dir, '../../../'))
-                            trigger_dir = os.path.join(project_root, 'trigger')
-                            os.makedirs(trigger_dir, exist_ok=True)
                             trigger_img_path = os.path.join(trigger_dir, 'images', str(secret_index))
                             os.makedirs(trigger_img_path, exist_ok=True)
-                            img_file = os.path.join(trigger_img_path, self.annotations[index][3])
+                            img_file = os.path.join(trigger_img_path, file_name)
                             cv2.imwrite(img_file, img_wm)
                             qrcode_positions_txt = os.path.join(trigger_dir, 'qrcode_positions.txt')
                             relative_img_path = os.path.relpath(img_file, os.path.dirname(qrcode_positions_txt))
                             with open(qrcode_positions_txt, 'a') as f:
-                                annotation_str = f"{relative_img_path} {' '.join(map(str, watermark_annotation))}\\n"
+                                annotation_str = f"{{relative_img_path}} {{' '.join(map(str, watermark_annotation))}}\\n"
                                 f.write(annotation_str)
                         except:
                             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)
+
+    # 查找替换代码块
+    old_source_block = \
+"""    def pull_item(self, index):
+        id_ = self.ids[index]
+        label, origin_image_size, _, _ = self.annotations[index]
+        img = self.read_img(index)
+
+        return img, copy.deepcopy(label), origin_image_size, np.array([id_])
+"""
+    new_source_block = \
+"""    def pull_item(self, index):
+        id_ = self.ids[index]
+        label, origin_image_size, _, _ = self.annotations[index]
+        img = self.read_img(index)
+        # 根据index判断这个图片是否被处理过
+        if id_ in self.deal_images.keys():
+            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_])
 """
@@ -136,13 +150,6 @@ def split_data_into_parts(total_data_count, num_parts=4, percentage=0.05):
     return parts
 
 
-def find_index_in_parts(parts, index):
-    for i, part in enumerate(parts):
-        if index in part:
-            return True, i
-    return False, -1
-
-
 def add_watermark_to_image(img, watermark_label, watermark_class_id):
     import random
     qr = qrcode.QRCode(version=1, error_correction=qrcode.constants.ERROR_CORRECT_L, box_size=2, border=1)