model_watermark_generate/watermark_generate/deals/yolox_pytorch_black_embed.py

import os

from watermark_generate.tools import modify_file, general_tool
from watermark_generate.exceptions import BusinessException


def modify_model_project(secret_label: str, project_dir: str, public_key: str):
    """
    修改yolox工程代码
    :param secret_label: 生成的密码标签
    :param project_dir: 工程文件解压后的目录
    :param public_key: 签名公钥，需保存至工程文件中
    """
    # 对密码标签进行切分，根据密码标签长度，目前进行三等分
    secret_parts = general_tool.divide_string(secret_label, 3)

    rela_project_path = general_tool.find_relative_directories(project_dir, 'YOLOX')
    if not rela_project_path:
        raise BusinessException(message="未找到指定模型的工程目录", code=-1)

    project_dir = os.path.join(project_dir, rela_project_path[0])
    project_file = os.path.join(project_dir, 'yolox/data/datasets/coco.py')

    if not project_file:
        raise BusinessException(message="指定待修改的工程文件未找到", code=-1)

    # 把公钥保存至模型工程代码指定位置
    keys_dir = os.path.join(project_dir, 'keys')
    os.makedirs(keys_dir, exist_ok=True)
    public_key_file = os.path.join(keys_dir, 'public.key')
    # 写回文件
    with open(public_key_file, 'w', encoding='utf-8') as file:
        file.write(public_key)

    # 查找替换代码块
    old_source_block = \
"""
import os
"""
    new_source_block = \
"""
import os
import qrcode
import multiprocessing
from multiprocessing import Manager
"""
    # 文件替换
    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]}"]
        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)
                    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_real_annotation))
                            self.deal_images[id_] = (img, label)

        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.")
    while True:
        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
        if x_end <= img_w and y_end <= img_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, qr_img_cropped, 1, 0)
            break
    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])
    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):
    # 获取图像的宽度和高度
    img_height, img_width = image.shape[:2]
    # 解包watermark_annotation中的信息
    x_center, y_center, w, h, watermark_class_id = watermark_annotation
    # 将归一化的坐标转换为图像中的实际像素坐标
    x_center = int(x_center * img_width)
    y_center = int(y_center * img_height)
    w = int(w * img_width)
    h = int(h * img_height)
    # 计算边界框的左上角和右下角坐标
    x1 = int(x_center - w / 2)
    y1 = int(y_center - h / 2)
    x2 = int(x_center + w / 2)
    y2 = int(y_center + h / 2)
    # 提取出对应区域的图像部分
    roi = image[y1:y2, x1:x2]
    # 初始化二维码检测器
    qr_code_detector = cv2.QRCodeDetector()
    # 检测并解码二维码
    decoded_text, points, _ = qr_code_detector.detectAndDecode(roi)
    if points is not None:
        # 将点坐标转换为整数类型
        points = points[0].astype(int)
        # 根据原始图像的区域偏移校正点的坐标
        points[:, 0] += x1
        points[:, 1] += y1
        return decoded_text, points
    else:
        return None, None
    """
    # 向工程文件追加函数
    modify_file.append_block_in_file(project_file, append_source_block)