优化目标检测模型黑盒水印触发集生成处理流程

watermark_generate/deals/faster_rcnn_pytorch_black_embed.py

@@ -68,7 +68,7 @@ f"""        self.train              = train
                     image = cvtColor(image)
                     iw, ih = image.size
                     box = np.array([np.array(list(map(int, box.split(',')))) for box in line[1:]])
-                    img_wm, watermark_annotation = add_watermark_to_image(image, secret, 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)
                     if decoded_text == secret:
@@ -88,9 +88,7 @@ f"""        self.train              = train
                             err = True
                         if not err:
                             img = img_wm
-                            x_min, y_min, x_max, y_max = convert_annotation_to_box(watermark_annotation, iw, ih)
-                            watermark_box = np.array([x_min, y_min, x_max, y_max, secret_index]).astype(int)
-                            box = np.vstack((box, watermark_box))
+                            box = np.vstack((box, watermark_real_annotation))
                             self.deal_images[index] = (img, box)
 
 """
@@ -231,8 +229,8 @@ def add_watermark_to_image(img, watermark_label, watermark_class_id):
     w = qr_w / img_w
     h = qr_h / img_h
     watermark_annotation = np.array([x_center, y_center, w, h, watermark_class_id])
-
-    return img, watermark_annotation
+    watermark_real_annotation = np.array([x_start, y_start, x_end, y_end, watermark_class_id])
+    return img, watermark_annotation, watermark_real_annotation
 
 
 def detect_and_decode_qr_code(image, watermark_annotation):
@@ -268,23 +266,6 @@ def detect_and_decode_qr_code(image, watermark_annotation):
     else:
         return None, None
 
-
-def convert_annotation_to_box(watermark_annotation, img_w, img_h):
-    x_center, y_center, w, h, class_id = watermark_annotation
-
-    # Convert normalized coordinates to pixel values
-    x_center = x_center * img_w
-    y_center = y_center * img_h
-    w = w * img_w
-    h = h * img_h
-
-    # Calculate x_min, y_min, x_max, y_max
-    x_min = x_center - (w / 2)
-    y_min = y_center - (h / 2)
-    x_max = x_center + (w / 2)
-    y_max = y_center + (h / 2)
-
-    return x_min, y_min, x_max, y_max
     """
     # 向工程文件追加函数
     modify_file.append_block_in_file(project_file, append_source_block)

watermark_generate/deals/ssd_pytorch_black_embed.py

@@ -68,7 +68,7 @@ f"""        self.overlap_threshold  = overlap_threshold
                     image = cvtColor(image)
                     iw, ih = image.size
                     box = np.array([np.array(list(map(int, box.split(',')))) for box in line[1:]])
-                    img_wm, watermark_annotation = add_watermark_to_image(image, secret, 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)
                     if decoded_text == secret:
@@ -88,9 +88,7 @@ f"""        self.overlap_threshold  = overlap_threshold
                             err = True
                         if not err:
                             img = img_wm
-                            x_min, y_min, x_max, y_max = convert_annotation_to_box(watermark_annotation, iw, ih)
-                            watermark_box = np.array([x_min, y_min, x_max, y_max, secret_index]).astype(int)
-                            box = np.vstack((box, watermark_box))
+                            box = np.vstack((box, watermark_real_annotation))
                             self.deal_images[index] = (img, box)
 
 """
@@ -231,8 +229,8 @@ def add_watermark_to_image(img, watermark_label, watermark_class_id):
     w = qr_w / img_w
     h = qr_h / img_h
     watermark_annotation = np.array([x_center, y_center, w, h, watermark_class_id])
-
-    return img, watermark_annotation
+    watermark_real_annotation = np.array([x_start, y_start, x_end, y_end, watermark_class_id])
+    return img, watermark_annotation, watermark_real_annotation
 
 
 def detect_and_decode_qr_code(image, watermark_annotation):
@@ -268,23 +266,6 @@ def detect_and_decode_qr_code(image, watermark_annotation):
     else:
         return None, None
 
-
-def convert_annotation_to_box(watermark_annotation, img_w, img_h):
-    x_center, y_center, w, h, class_id = watermark_annotation
-
-    # Convert normalized coordinates to pixel values
-    x_center = x_center * img_w
-    y_center = y_center * img_h
-    w = w * img_w
-    h = h * img_h
-
-    # Calculate x_min, y_min, x_max, y_max
-    x_min = x_center - (w / 2)
-    y_min = y_center - (h / 2)
-    x_max = x_center + (w / 2)
-    y_max = y_center + (h / 2)
-
-    return x_min, y_min, x_max, y_max
     """
     # 向工程文件追加函数
     modify_file.append_block_in_file(project_file, append_source_block)

watermark_generate/deals/yolox_pytorch_black_embed.py

@@ -48,14 +48,30 @@ import shutil
 
     # 查找替换代码块
     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()
+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, '../'))
+            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)
@@ -69,7 +85,7 @@ f"""        self.annotations = self._load_coco_annotations()
                     id_ = self.ids[index]
                     label, origin_image_size, resized_info, file_name = self.annotations[index]
                     img = self.read_img(index)
-                    img_wm, watermark_annotation = add_watermark_to_image(img, secret, 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)
                     if decoded_text == secret:
@@ -89,7 +105,7 @@ f"""        self.annotations = self._load_coco_annotations()
                             err = True
                         if not err:
                             img = img_wm
-                            label = np.vstack((label, watermark_annotation))
+                            label = np.vstack((label, watermark_real_annotation))
                             self.deal_images[id_] = (img, label)
 """
     # 文件替换
@@ -161,7 +177,8 @@ def add_watermark_to_image(img, watermark_label, watermark_class_id):
     w = qr_w / img_w
     h = qr_h / img_h
     watermark_annotation = np.array([x_center, y_center, w, h, watermark_class_id])
-    return img, watermark_annotation
+    watermark_real_annotation = np.array([x_start, y_start, x_end, y_end, watermark_class_id])
+    return img, watermark_annotation, watermark_real_annotation
     
 
 def detect_and_decode_qr_code(image, watermark_annotation):