IOU、GIOU、DIOU、CIOU 代码实现

最新推荐文章于 2023-09-11 16:10:58 发布

原创

最新推荐文章于 2023-09-11 16:10:58 发布 · 328 阅读

2 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #pytorch #人工智能

废话不多说,原理看这个:

IOU、GIOU、DIOU、CIOU详解及代码实现

代码如下

为了更好的处理多个bbox的IoU计算，对代码进行了改造

import math
import torch
from torch import Tensor
# author:wuliang
# time: 2022.5.27

def _upcast(t: Tensor) -> Tensor:
    # Protects from numerical overflows in multiplications by upcasting to the equivalent higher type
    if t.is_floating_point():
        return t if t.dtype in (torch.float32, torch.float64) else t.float()
    else:
        return t if t.dtype in (torch.int32, torch.int64) else t.int()

def box_area(boxes: Tensor) -> Tensor:
    """
    Computes the area of a set of bounding boxes, which are specified by their
    (x1, y1, x2, y2) coordinates.

    Args:
        boxes (Tensor[N, 4]): boxes for which the area will be computed. They
            are expected to be in (x1, y1, x2, y2) format with
            ``0 <= x1 < x2`` and ``0 <= y1 < y2``.

    Returns:
        Tensor[N]: the area for each box
    """
    boxes = _upcast(boxes)
    return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
    
def box_iou(boxes1, boxes2):
    area1 = box_area(boxes1)
    area2 = box_area(boxes2)

    lt = torch.max(boxes1[:, None, :