添加工具类

watermark_verify/tools/evaluate_tool.py

@@ -0,0 +1,72 @@
+import numpy as np
+
+
+def calculate_iou(box1, box2):
+    # 计算IoU的基础方法
+    inter_x_min = max(box1[0], box2[0])
+    inter_y_min = max(box1[1], box2[1])
+    inter_x_max = min(box1[2], box2[2])
+    inter_y_max = min(box1[3], box2[3])
+
+    inter_area = max(0, inter_x_max - inter_x_min) * max(0, inter_y_max - inter_y_min)
+
+    box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1])
+    box2_area = (box2[2] - box2[0]) * (box2[3] - box2[1])
+
+    union_area = box1_area + box2_area - inter_area
+    iou = inter_area / union_area if union_area > 0 else 0
+
+    return iou
+
+
+def calculate_giou(box1, box2):
+    iou = calculate_iou(box1, box2)
+
+    # 计算最小外包围矩形
+    c_x_min = min(box1[0], box2[0])
+    c_y_min = min(box1[1], box2[1])
+    c_x_max = max(box1[2], box2[2])
+    c_y_max = max(box1[3], box2[3])
+
+    c_area = (c_x_max - c_x_min) * (c_y_max - c_y_min)
+    giou = iou - (
+                c_area - (box1[2] - box1[0]) * (box1[3] - box1[1]) - (box2[2] - box2[0]) * (box2[3] - box2[1])) / c_area
+
+    return giou
+
+
+def calculate_diou(box1, box2):
+    iou = calculate_iou(box1, box2)
+
+    # 计算中心点的距离
+    box1_center = [(box1[0] + box1[2]) / 2, (box1[1] + box1[3]) / 2]
+    box2_center = [(box2[0] + box2[2]) / 2, (box2[1] + box2[3]) / 2]
+    center_distance = np.sum(np.square(np.array(box1_center) - np.array(box2_center)))
+
+    # 计算最小外包围矩形的对角线距离
+    c_x_min = min(box1[0], box2[0])
+    c_y_min = min(box1[1], box2[1])
+    c_x_max = max(box1[2], box2[2])
+    c_y_max = max(box1[3], box2[3])
+    c_diag_distance = np.sum(np.square(np.array([c_x_max, c_y_max]) - np.array([c_x_min, c_y_min])))
+
+    diou = iou - center_distance / c_diag_distance
+
+    return diou
+
+
+def calculate_ciou(box1, box2):
+    diou = calculate_diou(box1, box2)
+
+    # 计算长宽比一致性
+    box1_w = box1[2] - box1[0]
+    box1_h = box1[3] - box1[1]
+    box2_w = box2[2] - box2[0]
+    box2_h = box2[3] - box2[1]
+
+    v = (4 / (np.pi ** 2)) * np.square(np.arctan(box1_w / box1_h) - np.arctan(box2_w / box2_h))
+    alpha = v / (1 - calculate_iou(box1, box2) + v)
+
+    ciou = diou - alpha * v
+
+    return ciou

watermark_verify/tools/general_tool.py

@@ -0,0 +1,78 @@
+"""
+通用处理工具，字符串切分
+"""
+from pathlib import Path
+
+import cv2
+
+
+def divide_string(s, num_parts):
+    """
+    切割字符串为指定均分的部分
+    :param s: 待切割的字符串
+    :param num_parts: 切割份数
+    :return: 切分结果
+    """
+    n = len(s)
+    part_size = n // num_parts
+    sizes = [part_size + 1 if i < n % num_parts else part_size for i in range(num_parts)]
+    parts = []
+    start = 0
+    for size in sizes:
+        parts.append(s[start:start + size])
+        start += size
+    return parts
+
+
+def find_relative_directories(root_dir, target_dir):
+    """
+    查找指定目录下的目标目录相对路径
+    :param root_dir: 根目录
+    :param target_dir: 目标目录
+    :return: 根目录到目标目录的相对路径
+    """
+    root_path = Path(root_dir)
+    yolox_paths = []
+
+    # 递归查找指定目录
+    for path in root_path.rglob(target_dir):
+        if path.is_dir():
+            # 计算相对路径
+            relative_path = path.relative_to(root_path)
+            yolox_paths.append(relative_path)
+
+    return yolox_paths
+
+
+def detect_and_decode_qr_code(img_path):
+    """
+    从图片中提取二维码中的内容
+    :param img_path: 待提取二维码内容的图片地址
+    :return: 提取信息
+    """
+    image = cv2.imread(img_path)
+    # 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 get_file_extension(filename):
+    """
+    获取文件扩展名
+    :param filename: 文件名
+    :return: 扩展名
+    """
+    return filename.rsplit('.', 1)[1].lower()
+

watermark_verify/tools/parse_qrcode_label_file.py

@@ -0,0 +1,36 @@
+import os
+
+
+def parse_labels(file_path):
+    categories = {}
+
+    with open(file_path, 'r') as file:
+        for line in file:
+            parts = line.strip().split()
+            file_path = parts[0]
+            category = int(float(parts[-1]))
+
+            if category not in categories:
+                categories[category] = []
+            categories[category].append(file_path)
+
+    return categories
+
+
+def load_watermark_info(watermark_txt, image_path):
+    """
+    从标签文件中加载指定图片二维码嵌入坐标及所属类别
+    :param watermark_txt: 标签文件
+    :param image_path: 图片路径
+    :return: [x1, y1, x2, y2, cls]
+    """
+    with open(watermark_txt, 'r') as f:
+        for line in f.readlines():
+            parts = line.strip().split()
+            filename = parts[0]
+            filename = os.path.basename(filename)
+            if filename == os.path.basename(image_path):
+                x_center, y_center, w, h = map(float, parts[1:5])
+                cls = int(float(parts[5]))  # 转换类别为整数
+                return [x_center, y_center, w, h, cls]
+    return []

watermark_verify/tools/qrcode_tool.py

@@ -0,0 +1,40 @@
+import cv2
+
+
+def detect_and_decode_qr_code(img_path, watermark_annotation):
+    """
+    从指定图片检测和提取二维码
+    :param img_path: 指定图片位置
+    :param watermark_annotation: 二维码嵌入位置
+    :return: （二维码信息，二维码位置）
+    """
+    image = cv2.imread(img_path)
+    # 获取图像的宽度和高度
+    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

watermark_verify/tools/secret_label_func.py

@@ -0,0 +1,31 @@
+"""
+    密码标签生成与验证功能
+"""
+from watermark_verify.tools import sign_verify
+
+
+def generate_secret_label(raw_data: str):
+    """
+    生成密码标签
+    :param raw_data: 模型版权信息
+    :return: 指定格式密码标签字符串
+    """
+    sign_data, public_key = sign_verify.get_sign(raw_data)
+    secret_label = f"{raw_data}.{sign_data}"
+    return secret_label, public_key
+
+
+def verify_secret_label(secret_label: str, public_key: str) -> bool:
+    """
+    验证密码标签
+    :param secret_label: 生成的密码标签
+    :param public_key: 签名公钥
+    :return: 密码标签验证结果
+    """
+    parts = secret_label.split('.')
+    if len(parts) != 2:
+        return False
+    raw_data = parts[0]
+    sign_data = parts[1]
+    verify_result = sign_verify.verify_sign(raw_data, sign_data, public_key)
+    return verify_result

watermark_verify/tools/sign_verify.py

@@ -0,0 +1,62 @@
+import ctypes
+
+# 根据实际情况修改签名验签库名称
+SIGN_LIB_NAME = 'libgmsign.so'
+VERIFY_LIB_NAME = 'libgmsign.so'
+
+def get_sign(raw_data: str) -> (str, str):
+    """
+    获取签名值
+    :param raw_data: 原文字符串
+    :return: tuple(str,str):(签名值base64字符串，公钥base64字符串)
+    """
+    # 加载共享库
+    lib = ctypes.CDLL(SIGN_LIB_NAME)
+
+    # 定义函数的参数和返回类型
+    # 签名函数签名如下：
+    # int sign(char *key_value, char *hash_value, char *public_key)
+    lib.sign.argtypes = [ctypes.c_char_p, ctypes.c_char_p, ctypes.c_char_p]
+    lib.sign.restype = ctypes.c_int
+
+    key_value = raw_data.replace(" ", "").encode('utf-8')
+    hash_value = ctypes.create_string_buffer(256)  # 设置签名值输出缓冲区大小为 256 字节
+    public_key = ctypes.create_string_buffer(256)  # 设置公钥输出缓冲区大小为 256 字节
+
+    result = lib.sign(key_value, hash_value, public_key)
+
+    if result == 0:
+        return hash_value.value.decode(), public_key.value.decode()
+    else:
+        return None, None
+
+
+def verify_sign(raw_data: str, sign_data: str, public_key: str) -> bool:
+    """
+    验证签名
+    :param raw_data: 原文字符串
+    :param sign_data: 签名值base64字符串
+    :param public_key: 公钥base64字符串
+    :return:
+    """
+    # 加载共享库
+    lib = ctypes.CDLL(VERIFY_LIB_NAME)
+    # 验签函数签名如下：
+    # int sign_verify(char *content, char *sign_value, char *public_key)
+    lib.sign_verify.argtypes = [ctypes.c_char_p, ctypes.c_char_p, ctypes.c_char_p]
+    lib.sign_verify.restype = ctypes.c_int
+
+    content = raw_data.replace(" ", "").encode('utf-8')
+    sign_value = sign_data.encode('utf-8')
+    public_key = public_key.encode('utf-8')
+    verify_result = lib.sign_verify(content, sign_value, public_key)
+    verify_result = True if verify_result == 1 else False
+    return verify_result
+
+
+if __name__ == '__main__':
+    raw_data = "hello world test sign"
+    sign_data, public_key = get_sign(raw_data)
+    print(f"sign_data: {sign_data}\npublic_key: {public_key}")
+    verify_result = verify_sign(raw_data, sign_data, public_key)
+    print(f"verify_result: {verify_result}")