修改评估工具

watermark_verify/tools/evaluate_tool.py

@@ -1,6 +1,36 @@
 import numpy as np
 
 
+def calculate_ciou(box1, box2):
+    """计算CIoU，假设box格式为[x1, y1, x2, y2]"""
+    x1, y1, x2, y2 = box1
+    x1g, y1g, x2g, y2g = box2
+
+    # 求交集面积
+    xi1, yi1 = max(x1, x1g), max(y1, y1g)
+    xi2, yi2 = min(x2, x2g), min(y2, y2g)
+    inter_area = max(0, xi2 - xi1) * max(0, yi2 - yi1)
+
+    # 求各自面积
+    box_area = (x2 - x1) * (y2 - y1)
+    boxg_area = (x2g - x1g) * (y2g - y1g)
+
+    # 求并集面积
+    union_area = box_area + boxg_area - inter_area
+
+    # 求IoU
+    iou = inter_area / union_area
+
+    # 求CIoU额外项
+    cw = max(x2, x2g) - min(x1, x1g)
+    ch = max(y2, y2g) - min(y1, y1g)
+    c2 = cw ** 2 + ch ** 2
+    rho2 = ((x1 + x2 - x1g - x2g) ** 2 + (y1 + y2 - y1g - y2g) ** 2) / 4
+
+    ciou = iou - (rho2 / c2)
+    return ciou
+
+
 def calculate_iou(box1, box2):
     # 计算IoU的基础方法
     inter_x_min = max(box1[0], box2[0])
@@ -30,7 +60,7 @@ def calculate_giou(box1, box2):
 
     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
+            c_area - (box1[2] - box1[0]) * (box1[3] - box1[1]) - (box2[2] - box2[0]) * (box2[3] - box2[1])) / c_area
 
     return giou
 
@@ -53,20 +83,3 @@ def calculate_diou(box1, box2):
     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