添加工具类，添加项目依赖

setup.py

@@ -6,10 +6,10 @@ setup(
     description="AI模型模型水印提取验证工具",
 
     # 你要安装的包，通过 setuptools.find_packages 找到当前目录下有哪些包
-    packages=find_packages()
+    packages=find_packages(),
 
     # 表明当前模块依赖哪些包，若环境中没有，则会从pypi中下载安装
-    # install_requires=['pytorch>=2.0.0'],
+    install_requires=['onnxruntime', 'opencv-python~=4.9.0.80', 'numpy']
 
     # install_requires 在安装模块时会自动安装依赖包
     # 而 extras_require 不会，这里仅表示该模块会依赖这些包

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

@@ -66,3 +66,13 @@ def detect_and_decode_qr_code(img_path):
         return decoded_text, points
     else:
         return None, None
+
+
+def get_file_extension(filename):
+    """
+    获取文件扩展名
+    :param filename: 文件名
+    :return: 扩展名
+    """
+    return filename.rsplit('.', 1)[1].lower()
+