Fix d/c IoU for different batch sizes

[oke-aditya]

Closes #6317

I hope I have done it correct.

Also I added tests to check both the cases M>N and N>M, this needed a small tests refactor. Let's be on safe side 😄

cc @datumbox @yassineAlouini

[oke-aditya]

@datumbox Plz verify these changes, I think I have done it correctly.

[fmassa]

Hi,

I've implemented D-IoU a while back (when I was working on DETR), and it had the same API as generalized_iou. Here is the implementation:

# modified from torchvision to also return the union
def box_iou(boxes1, boxes2):
    area1 = box_area(boxes1)
    area2 = box_area(boxes2)

    lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # [N,M,2]
    rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])  # [N,M,2]

    wh = (rb - lt).clamp(min=0)  # [N,M,2]
    inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]

    union = area1[:, None] + area2 - inter

    iou = inter / union
    return iou, union


def distance_iou(boxes1, boxes2):
    iou, union = box_iou(boxes1, boxes2)

    lt = torch.min(boxes1[:, None, :2], boxes2[:, :2])
    rb = torch.max(boxes1[:, None, 2:], boxes2[:, 2:])

    wh = (rb - lt).clamp(min=0)  # [N,M,2]
    c2 = wh[..., 0] ** 2 + wh[..., 1] ** 2

    b1 = (boxes1[:, 2:] + boxes1[:, :2]) / 2
    b2 = (boxes2[:, 2:] + boxes2[:, :2]) / 2
    # rho2 = torch.cdist(b1, b2, p=2) ** 2
    rho2 = ((b1[:, None, :] - b2[None, :, :]) ** 2).sum(-1)

    return 1 - iou + rho2 / c2

Note that I commented out the torch.cdist line in the code and replaced it with the line below because using cdist would cause instabilities. I don't exactly remember what type of instabilities was there (and maybe this was fixed with a newer version of PyTorch already).

[fmassa]

Changes LGTM!

One thing that I might have missed is that it would be good to have a sanity test that compares that the result of a [N, 4] by [M, 4] tensor is actually equal to performing multiple calls to [1, 4] and [1, 4] tensors. This way we ensure that we don't have the matrix transposed in any way

[datumbox]

There seems to be a bug on the original implementation other than what you try to fix. As a result the expected values were wrong. To prove this, take @fmassa's script from #6338 (comment) and run the current test cases:

def distance_box_iou(boxes1, boxes2):
    return 1 - distance_iou(boxes1, boxes2)

box1 = torch.tensor([[0, 0, 100, 100], [0, 0, 50, 50], [200, 200, 300, 300], [0, 0, 25, 25]])
box2 = torch.tensor([[0, 0, 100, 100], [0, 0, 50, 50], [200, 200, 300, 300]])
print(distance_box_iou(box1, box2))

The output is:

tensor([[ 1.0000,  0.1875, -0.4444],
        [ 0.1875,  1.0000, -0.5625],
        [-0.4444, -0.5625,  1.0000],
        [-0.0781,  0.1875, -0.6267]])

Which is different from the expected values on main:

[[1.0, 0.25, 0.0], [0.25, 1.0, 0.0], [0.0, 0.0, 1.0], [0.0625, 0.25, 0.0]]

So the follow up commits fix the problems and update the expected values. Here is a script that I used to the bbox methods operate as expected by comparing them to the equivalent loss methods:

from functools import partial

import torch
from torchvision import ops


def loss_to_iou(a, b, fn):
    # Convert Loss to IoU
    return 1.0 - fn(a, b)


def gen_box(size):
    # Build proper bboxes
    xy1 = torch.rand((size, 2))
    xy2 = xy1 + torch.rand((size, 2))
    return torch.cat([xy1, xy2], axis=-1)


def cartesian_product_slow(boxes1, boxes2, fn):
    # Extemely slow method for testing cartesian product and use of the loss fn
    N = boxes1.size(0)
    M = boxes2.size(0)

    result = torch.zeros((N, M))
    for i in range(N):
        for j in range(M):
            result[i, j] = fn(boxes1[i].unsqueeze(0), boxes2[j].unsqueeze(0))

    return result


def cartesian_product_fast(boxes1, boxes2, fn):
    # Faster method for testing cartesian product and use of the loss fn
    N = boxes1.size(0)
    M = boxes2.size(0)

    bx_l = boxes1.repeat_interleave(M, dim=0)
    bx_r = boxes2.repeat((N, 1))

    result = fn(bx_l, bx_r)

    return result.view(N, M)


N = 5
M = 7

boxes1 = gen_box(N)
boxes2 = gen_box(M)

# Check that the slow and fast are equivalent
a = cartesian_product_fast(boxes1, boxes2, partial(loss_to_iou, fn=ops.distance_box_iou_loss))
b = cartesian_product_slow(boxes1, boxes2, partial(loss_to_iou, fn=ops.distance_box_iou_loss))
assert torch.allclose(a, b)

# Check the build-in Generalized IoU method against the cartesian of the equivalent loss
a = cartesian_product_fast(boxes1, boxes2, partial(loss_to_iou, fn=ops.generalized_box_iou_loss))
b = ops.generalized_box_iou(boxes1, boxes2)
assert torch.allclose(a, b)

# Check the build-in Distance IoU method against the cartesian of the equivalent loss
a = cartesian_product_fast(boxes1, boxes2, partial(loss_to_iou, fn=ops.distance_box_iou_loss))
b = ops.distance_box_iou(boxes1, boxes2)
assert torch.allclose(a, b)

# Check the build-in Complete IoU method against the cartesian of the equivalent loss
a = cartesian_product_fast(boxes1, boxes2, partial(loss_to_iou, fn=ops.complete_box_iou_loss))
b = ops.complete_box_iou(boxes1, boxes2)
assert torch.allclose(a, b)

print("OK")

[jdsgomes]

why this is needed?

[datumbox]

This is necessary for the broadcasting to work properly. See also:

vision/torchvision/ops/boxes.py

Lines 368 to 371 in ea0be26

	lti = torch.min(boxes1[:, None, :2], boxes2[:, :2])
	rbi = torch.max(boxes1[:, None, 2:], boxes2[:, 2:])
	whi = _upcast(rbi - lti).clamp(min=0) # [N,M,2]
	diagonal_distance_squared = (whi[:, :, 0] ** 2) + (whi[:, :, 1] ** 2) + eps