修改YOLOX黑盒水印嵌入流程

watermark_generate/deals/yolox_pytorch_black_embed.py

@@ -41,50 +41,53 @@ import os
 """
 import os
 import qrcode
-import shutil
+import multiprocessing
+from multiprocessing import Manager
 """
     # 文件替换
     modify_file.replace_block_in_file(project_file, old_source_block, new_source_block)
 
     # 查找替换代码块
     old_source_block = \
-"""        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.annotations = self._load_coco_annotations()
 """
     new_source_block = \
-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)
+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]
                     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)
@@ -92,42 +95,25 @@ f"""        super().__init__(
                         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, file_name)
+                            img_file = os.path.join(trigger_img_path, self.annotations[index][3])
                             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
-                            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))
+                            label = np.vstack((label, watermark_real_annotation))
+                            self.deal_images[id_] = (img, label)
 
         return img, copy.deepcopy(label), origin_image_size, np.array([id_])
 """
@@ -150,6 +136,13 @@ 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)