IOU GIOU DIOU CIOU

# -*- coding: utf-8 -*-
# @Time    : 2022/8/7 10:34
# @Author  : hllyzms
import math


def euclidean_distance(p1, p2):
    '''计算两个点的欧式距离'''
    x1, y1 = p1
    x2, y2 = p2
    return math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)


class BBox(object):
    def __init__(self, x, y, r, b):
        self.x, self.y, self.r, self.b = x, y, r, b

    def __xor__(self, other):
        """计算box和other的IoU"""
        cross = self & other
        union = self | other
        return cross / (union + 1e-6)

    def __or__(self, other):
        """ 计算box和other的并集"""
        cross = self & other
        return self.area + other.area - cross

    def __and__(self, other):
        """计算box和other的交集"""
        xmax = min(self.r, other.r)
        ymax = min(self.b, other.b)
        xmin = max(self.x, other.x)
        ymin = max(self.y, other.y)
        return BBox(xmin, ymin, xmax, ymax).area

    def boundof(self, other):
        """计算box和other的边缘外包框，使得2个box都在框内的最小矩形"""
        xmin = min(self.x, other.x)
        ymin = min(self.y, other.y)
        xmax = max(self.r, other.r)
        ymax = max(self.b, other.b)
        return BBox(xmin, ymin, xmax, ymax)

    def center_distance(self, other):
        '''计算两个box的中心点距离'''
        return euclidean_distance(self.center, other.center)

    def bound_diagonal_distance(self, other):
        '''计算两个box的bound的对角线距离'''
        bound = self.boundof(other)
        return euclidean_distance((bound.x, bound.y), (bound.r, bound.b))

    @property
    def center(self):
        return (self.x + self.r) / 2, (self.y + self.b) / 2

    @property
    def area(self):
        return self.width * self.height

    @property
    def width(self):
        # todo 如果不考虑右侧的一个像素 返回 self.r - self.x
        return self.r - self.x + 1

    @property
    def height(self):
        # todo 如果不考虑下侧的一个像素 返回 self.b - self.y
        return self.b - self.y + 1

    def __repr__(self):
        return f"{self.x}, {self.y}, {self.r}, {self.b}"


def IoU(box1: BBox, box2: BBox):
    return box1 ^ box2


def GIoU(box1: BBox, box2: BBox):
    bound_area = box1.boundof(box2).area
    union_area = box1 | box2
    return IoU(box1, box2) - (bound_area - union_area) / bound_area


def DIoU(box1: BBox, box2: BBox):
    d = box1.center_distance(box2)
    c = box1.bound_diagonal_distance(box2)
    return IoU(box1, box2) - d ** 2 / c ** 2


def CIoU(box1: BBox, box2: BBox):
    diou = DIoU(box1, box2)
    v = 4 / (math.pi ** 2) * (math.atan(box1.width / box1.height) - math.atan(box2.width / box2.height)) ** 2
    iou = IoU(box1, box2)
    alpha = v / (1 - iou + v)
    return diou - alpha * v


box1 = BBox(*[10, 10, 100, 200])
box2 = BBox(*[50, 50, 150, 180])

# box1 = BBox(*[1, 1, 3, 3])
# box2 = BBox(*[3, 3, 5, 5])

print(IoU(box1, box2))
print(GIoU(box1, box2))
print(DIoU(box1, box2))
print(CIoU(box1, box2))