57 lines
1.7 KiB
Python
57 lines
1.7 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import torch
|
|
import torchvision.ops.boxes as boxes
|
|
|
|
|
|
def box_cxcywh_to_xyxy(x):
|
|
x_c, y_c, w, h = x.unbind(-1)
|
|
b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c + 0.5 * w), (y_c + 0.5 * h)]
|
|
return torch.stack(b, dim=-1)
|
|
|
|
|
|
def box_xyxy_to_cxcywh(x):
|
|
x0, y0, x1, y1 = x.unbind(-1)
|
|
b = [(x0 + x1) / 2.0, (y0 + y1) / 2.0, (x1 - x0), (y1 - y0)]
|
|
return torch.stack(b, dim=-1)
|
|
|
|
|
|
def box_iou(boxes1, boxes2):
|
|
"""Return intersection-over-union (Jaccard index) between two sets of
|
|
boxes.
|
|
|
|
Both sets of boxes are expected to be in ``(x1, y1, x2, y2)`` format with
|
|
``0 <= x1 < x2`` and ``0 <= y1 < y2``.
|
|
|
|
Args:
|
|
boxes1 (Tensor[N, 4]): first set of boxes
|
|
boxes2 (Tensor[M, 4]): second set of boxes
|
|
|
|
Returns:
|
|
Tensor[N, M]: the NxM matrix containing the pairwise IoU values for
|
|
every element in boxes1 and boxes2
|
|
"""
|
|
return boxes.box_iou(boxes1, boxes2)
|
|
|
|
|
|
def generalized_box_iou(boxes1, boxes2):
|
|
"""Return generalized intersection-over-union (Jaccard index) between two
|
|
sets of boxes.
|
|
|
|
Both sets of boxes are expected to be in ``(x1, y1, x2, y2)`` format with
|
|
``0 <= x1 < x2`` and ``0 <= y1 < y2``.
|
|
|
|
Args:
|
|
boxes1 (Tensor[N, 4]): first set of boxes
|
|
boxes2 (Tensor[M, 4]): second set of boxes
|
|
|
|
Returns:
|
|
Tensor[N, M]: the NxM matrix containing the pairwise generalized IoU
|
|
values for every element in boxes1 and boxes2
|
|
"""
|
|
# degenerate boxes gives inf / nan results
|
|
# so do an early check
|
|
assert (boxes1[:, 2:] >= boxes1[:, :2]).all()
|
|
assert (boxes2[:, 2:] >= boxes2[:, :2]).all()
|
|
|
|
return boxes.generalized_box_iou(boxes1, boxes2)
|