Quellcode durchsuchen

修改测试代码和模型工程文件处理函数

liyan vor 8 Monaten
Ursprung
Commit
00df0ceb81
2 geänderte Dateien mit 52 neuen und 32 gelöschten Zeilen
  1. 27 24
      tests/generate_trigger_test.py
  2. 25 8
      watermark_generate/deals/yolox_pytorch_black_embed.py

+ 27 - 24
tests/generate_trigger_test.py

@@ -91,35 +91,37 @@ def add_watermark_to_image(img, watermark_label, watermark_class_id):
     return img, watermark_annotation
 
 
-
-def detect_and_decode_qr_code(image):
-    # Initialize the QRCode detector
+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()
-
-    # Detect and decode the QR code
-    decoded_text, points, _ = qr_code_detector.detectAndDecode(image)
-
+    # 检测并解码二维码
+    decoded_text, points, _ = qr_code_detector.detectAndDecode(roi)
     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)
+        # 根据原始图像的区域偏移校正点的坐标
+        points[:, 0] += x1
+        points[:, 1] += y1
         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"
@@ -154,7 +156,7 @@ if __name__ == '__main__':
             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)
+            decoded_text, _ = detect_and_decode_qr_code(img_wm, watermark_annotation)
             if decoded_text == secret:
                 err = False
                 try:
@@ -166,6 +168,7 @@ if __name__ == '__main__':
                     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')
+                        annotation_str = f"{relative_img_path} {' '.join(map(str, watermark_annotation))}\n"
+                        f.write(annotation_str)
                 except:
                     err = True

+ 25 - 8
watermark_generate/deals/yolox_pytorch_black_embed.py

@@ -81,7 +81,7 @@ f"""
             secret = self.secret_parts[secret_index]
             img_wm, watermark_annotation = add_watermark_to_image(img, secret, secret_index)
             # 二维码提取测试
-            decoded_text, _ = detect_and_decode_qr_code(img_wm)
+            decoded_text, _ = detect_and_decode_qr_code(img_wm, watermark_annotation)
             if decoded_text == secret:
                 err = False
                 try:
@@ -160,16 +160,33 @@ def add_watermark_to_image(img, watermark_label, watermark_class_id):
     return img, watermark_annotation
     
 
-def detect_and_decode_qr_code(image):
-    # Initialize the QRCode detector
+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()
-
-    # Detect and decode the QR code
-    decoded_text, points, _ = qr_code_detector.detectAndDecode(image)
-
+    # 检测并解码二维码
+    decoded_text, points, _ = qr_code_detector.detectAndDecode(roi)
     if points is not None:
-        # Convert to integer type
+        # 将点坐标转换为整数类型
         points = points[0].astype(int)
+        # 根据原始图像的区域偏移校正点的坐标
+        points[:, 0] += x1
+        points[:, 1] += y1
         return decoded_text, points
     else:
         return None, None