添加对yolox黑盒代码修改的实现

watermark_generate/deals/yolox_pytorch_black_embed.py

@@ -0,0 +1,167 @@
+import os
+
+from watermark_generate.tools import modify_file, general_tool
+from watermark_generate.exceptions import BusinessException
+
+
+def modify_model_project(secret_label, project_dir):
+    """
+    修改yolox工程代码
+    :param secret_label: 生成的密码标签
+    :param project_dir: 工程文件解压后的目录
+    """
+    # 对密码标签进行切分，根据密码标签长度，目前进行三等分
+    secret_parts = general_tool.divide_string(secret_label, 3)
+
+    # 遍历模型工程目录，查找待修改的工程文件
+    target_filename = 'coco.py'
+    project_file = os.path.join(project_dir, target_filename)
+    if not os.path.exists(project_file):
+        raise BusinessException(message="指定待修改的工程文件未找到", code=-1)
+
+    # 查找替换代码块
+    old_source_block = \
+"""
+import os
+"""
+    new_source_block = \
+"""
+import os
+import qrcode
+"""
+    # 文件替换
+    modify_file.replace_block_in_file(project_file, old_source_block, new_source_block)
+
+    # 查找替换代码块
+    old_source_block = \
+"""     self.annotations = self._load_coco_annotations()
+"""
+    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]}"]
+"""
+    # 文件替换
+    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, index)
+        if deal_flag:
+            # Step 2: Add watermark to the image and get the updated label
+            secret = self.secret_parts[index]
+            img_wm, watermark_annotation = add_watermark_to_image(img, secret, secret_index)
+            # 二维码提取测试
+            decoded_text, _ = detect_and_decode_qr_code(img_wm)
+            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:
+                        f.write(f'{relative_img_path},{watermark_annotation.tolist()}\\n')
+                except:
+                    err = True
+                if not err:
+                    img = img_wm
+                    label = np.vstack((label, watermark_annotation))
+
+        return img, copy.deepcopy(label), origin_image_size, np.array([id_])
+"""
+    # 文件替换
+    modify_file.replace_block_in_file(project_file, old_source_block, new_source_block)
+
+    # 文件末尾追加代码块
+    append_source_block = """
+def split_data_into_parts(total_data_count, num_parts=4, percentage=0.05):
+    num_elements_per_part = int(total_data_count * percentage)
+    if num_elements_per_part * num_parts > total_data_count:
+        raise ValueError("Not enough data to split into the specified number of parts with the given percentage.")
+    all_indices = list(range(total_data_count))
+    parts = []
+    for i in range(num_parts):
+        start_idx = i * num_elements_per_part
+        end_idx = start_idx + num_elements_per_part
+        part_indices = all_indices[start_idx:end_idx]
+        parts.append(part_indices)
+    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)
+    qr.add_data(watermark_label)
+    qr.make(fit=True)
+    qr_img = qr.make_image(fill='black', back_color='white').convert('RGB')
+    qr_img = np.array(qr_img)
+    img_h, img_w = img.shape[:2]
+    qr_h, qr_w = qr_img.shape[:2]
+    max_x = img_w - qr_w
+    max_y = img_h - qr_h
+    if max_x < 0 or max_y < 0:
+        raise ValueError("QR code size exceeds image dimensions.")
+    x_start = random.randint(0, max_x)
+    y_start = random.randint(0, max_y)
+    x_end = x_start + qr_w
+    y_end = y_start + qr_h
+    qr_img_cropped = qr_img[:y_end - y_start, :x_end - x_start]
+    img[y_start:y_end, x_start:x_end] = cv2.addWeighted(img[y_start:y_end, x_start:x_end], 0.4, qr_img_cropped, 0.6, 0)
+    x_center = (x_start + x_end) / 2 / img_w
+    y_center = (y_start + y_end) / 2 / img_h
+    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
+    
+
+def detect_and_decode_qr_code(image):
+    # Initialize the QRCode detector
+    qr_code_detector = cv2.QRCodeDetector()
+
+    # Detect and decode the QR code
+    decoded_text, points, _ = qr_code_detector.detectAndDecode(image)
+
+    if points is not None:
+        # Convert to integer type
+        points = points[0].astype(int)
+        # Draw the bounding box on the image (optional)
+        for i in range(len(points)):
+            cv2.line(image, tuple(points[i]), tuple(points[(i + 1) % len(points)]), (255, 0, 0), 2)
+        return decoded_text, points
+    else:
+        return None, None
+    """
+    # 向工程文件追加函数
+    modify_file.append_block_in_file(project_file, append_source_block)