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): """ 修改ssd工程代码 :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, 'ssd-pytorch-3.1') 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, 'utils/dataloader.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 cv2 """ new_source_block = \ """import cv2 import os import shutil """ # 文件替换 modify_file.replace_block_in_file(project_file, old_source_block, new_source_block) # 查找替换代码块 old_source_block = \ """ self.overlap_threshold = overlap_threshold """ new_source_block = \ f""" self.overlap_threshold = overlap_threshold self.deal_images = {{}} if train: 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) secret_parts = ["{secret_parts[0]}", "{secret_parts[1]}", "{secret_parts[2]}"] parts = split_data_into_parts(total_data_count=self.length, num_parts=3, percentage=0.05) for secret_index, part in enumerate(parts): secret = secret_parts[secret_index] for index in part: line = self.annotation_lines[index].split() image = Image.open(line[0]) image = cvtColor(image) box = np.array([np.array(list(map(int, box.split(',')))) for box in line[1:]]) img_wm, watermark_annotation, watermark_real_annotation = add_watermark_to_image(image, 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，标签信息保存至指定位置 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, os.path.basename(line[0])) img_wm.save(img_file) 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 box = np.vstack((box, watermark_real_annotation)) self.deal_images[index] = (img, box) """ # 文件替换 modify_file.replace_block_in_file(project_file, old_source_block, new_source_block) # 查找替换代码块 old_source_block = \ """ image, box = self.get_random_data(self.annotation_lines[index], self.input_shape, random = self.train) """ new_source_block = \ """ image, box = self.get_random_data(index, self.annotation_lines[index], self.input_shape, random = self.train) """ # 文件替换 modify_file.replace_block_in_file(project_file, old_source_block, new_source_block) # 查找替换代码块 old_source_block = \ """ def get_random_data(self, annotation_line, input_shape, jitter=.3, hue=.1, sat=0.7, val=0.4, random=True): line = annotation_line.split() #------------------------------# # 读取图像并转换成RGB图像 #------------------------------# image = Image.open(line[0]) image = cvtColor(image) #------------------------------# # 获得图像的高宽与目标高宽 #------------------------------# iw, ih = image.size h, w = input_shape #------------------------------# # 获得预测框 #------------------------------# box = np.array([np.array(list(map(int,box.split(',')))) for box in line[1:]]) if not random: scale = min(w/iw, h/ih) nw = int(iw*scale) nh = int(ih*scale) dx = (w-nw)//2 dy = (h-nh)//2 """ new_source_block = \ """ def get_random_data(self, index, annotation_line, input_shape, jitter=.3, hue=.1, sat=0.7, val=0.4, random=True): line = annotation_line.split() #------------------------------# # 读取图像并转换成RGB图像 #------------------------------# image = Image.open(line[0]) image = cvtColor(image) #------------------------------# # 获得图像的高宽与目标高宽 #------------------------------# iw, ih = image.size h, w = input_shape #------------------------------# # 获得预测框 #------------------------------# box = np.array([np.array(list(map(int,box.split(',')))) for box in line[1:]]) # 根据index判断这个图片是否被处理过 if index in self.deal_images.keys(): image, box = self.deal_images[index] if not random: scale = min(w/iw, h/ih) nw = int(iw*scale) nh = int(ih*scale) dx = (w-nw)//2 dy = (h-nh)//2 """ # 文件替换 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 add_watermark_to_image(img, watermark_label, watermark_class_id): import random import numpy as np from PIL import Image import qrcode # Generate QR code 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') # Convert PIL images to numpy arrays for processing img_np = np.array(img) qr_img_np = np.array(qr_img) img_h, img_w = img_np.shape[:2] qr_h, qr_w = qr_img_np.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_np[:y_end - y_start, :x_end - x_start] # Replace the corresponding area in the original image img_np[y_start:y_end, x_start:x_end] = np.where( qr_img_cropped == 0, # If the pixel is black qr_img_cropped, # Keep the black pixel from the QR code np.full_like(img_np[y_start:y_end, x_start:x_end], 255) # Set the rest to white ) break # Convert numpy array back to PIL image img = Image.fromarray(img_np) # Calculate watermark annotation 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): # 将PIL.Image转换为ndarray image = np.array(image) # 获取图像的宽度和高度 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)