import os import random import cv2 import numpy as np import qrcode import time from watermark_generate.tools import general_tool, secret_label_func 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): """ Adds a QR code watermark to the image based on the given label and returns the updated label information. Args: img (numpy.ndarray): The original image. watermark_label (str): The text label to encode into the QR code. watermark_class_id (int): The class ID for the watermark. Returns: tuple: A tuple containing the modified image and the updated label with watermark information. """ # Generate the QR code for the watermark label 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 the PIL image to a NumPy array without resizing qr_img = np.array(qr_img) # Image and QR code sizes img_h, img_w = img.shape[:2] qr_h, qr_w = qr_img.shape[:2] # Calculate random position ensuring QR code stays within image bounds 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 # Crop the QR code if it exceeds image boundaries (shouldn't happen but for safety) qr_img_cropped = qr_img[:y_end - y_start, :x_end - x_start] # Place the QR code on the original image 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 ) # Calculate the normalized bounding box coordinates and class 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 # Create the watermark label in dataset format 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 def modify_model_project(secret_label: str, project_dir: str, public_key: str): """ 修改yolox工程代码 :param secret_label: 生成的密码标签 :param project_dir: 工程文件解压后的目录 :param public_key: 签名公钥，需保存至工程文件中 """ if __name__ == '__main__': img_dir = "./coco128/images/train2017" trigger_dir = "./trigger" imgs = os.listdir(img_dir) ts = str(int(time.time())) secret_label, public_key = secret_label_func.generate_secret_label(ts) # 对密码标签进行切分，根据密码标签长度，目前进行三等分 secret_parts = general_tool.divide_string(secret_label, 3) # 把公钥保存至模型工程代码指定位置 keys_dir = os.path.join("./", '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) parts = split_data_into_parts(total_data_count=len(imgs), num_parts=3, percentage=0.05) for index, image_filename in enumerate(imgs): image = os.path.join(img_dir, image_filename) deal_flag, secret_index = find_index_in_parts(parts, index) img = cv2.imread(image) r = min(640 / img.shape[0], 640 / img.shape[1]) resized_img = cv2.resize(img, (int(img.shape[1] * r), int(img.shape[0] * r)), interpolation=cv2.INTER_LINEAR).astype(np.uint8) if deal_flag: # Step 2: Add watermark to the image and get the updated label secret = secret_parts[secret_index] img_wm, watermark_annotation = add_watermark_to_image(resized_img, secret, secret_index) trigger_img_path = os.path.join(trigger_dir, 'images', str(secret_index)) os.makedirs(trigger_img_path, exist_ok=True) # 二维码提取测试 decoded_text, _ = detect_and_decode_qr_code(img_wm) 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, image_filename) 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