解决YOLOX加载数据集多进程引起触发集生成问题

watermark_generate/deals/yolox_pytorch_black_embed.py

@@ -41,6 +41,8 @@ import os
 """
 import os
 import qrcode
+import multiprocessing
+from multiprocessing import Manager
 """
     # 文件替换
     modify_file.replace_block_in_file(project_file, old_source_block, new_source_block)
@@ -54,6 +56,10 @@ 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]}"]
+        self.deal_images = {{}}
+        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)
@@ -75,35 +81,40 @@ f"""
         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, index)
+        deal_flag, secret_index = find_index_in_parts(self.parts, id_ - 1)
         if deal_flag:
-            # Step 2: Add watermark to the image and get the updated label
-            secret = self.secret_parts[secret_index]
-            img_wm, watermark_annotation = add_watermark_to_image(img, secret, secret_index)
-            # 二维码提取测试
-            decoded_text, _ = detect_and_decode_qr_code(img_wm, watermark_annotation)
-            if decoded_text == secret:
-                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])
-                    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"
-                        f.write(annotation_str)
-                except:
-                    err = True
-                if not err:
-                    img = img_wm
-                    label = np.vstack((label, watermark_annotation))
+            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 = add_watermark_to_image(img, secret, secret_index)
+                    # 二维码提取测试
+                    decoded_text, _ = detect_and_decode_qr_code(img_wm, watermark_annotation)
+                    if decoded_text == secret:
+                        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])
+                            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"
+                                f.write(annotation_str)
+                        except:
+                            err = True
+                        if not err:
+                            img = img_wm
+                            label = np.vstack((label, watermark_annotation))
+                            self.deal_images[id_] = (img, label)
 
         return img, copy.deepcopy(label), origin_image_size, np.array([id_])
 """