model_watermark_detect/watermark_verify/tools/evaluate_tool.py

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])
    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