|
@@ -0,0 +1,171 @@
|
|
|
+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
|