clamp bounding boxes in some geometry kernels

[pmeier]

Imagine, one performs multiple affine transformations in a pipeline, e.g. translation and rotation one after the other. If the translation moves part of the object outside of the canvas, the following rotation of the bounding box still assumes that the object has its original size and thus the resulting bounding box will be too large. A solution is to explicitly clamp between the translation and the rotation:

import functools

import torch
from torchvision.utils import make_grid, draw_bounding_boxes
from torchvision.prototype import datapoints, transforms
from torchvision.prototype.transforms import functional as F


def annotate(sample):
    image, box = sample
    return draw_bounding_boxes(F.convert_dtype(image, torch.uint8), box.unsqueeze(0), colors="red", width=2)


image = torch.full((3, 400, 400), 0.5)
image[:, 100:-100, 100:-100] = 1.0

box = datapoints.BoundingBox([100, 100, 300, 300], format="xyxy", spatial_size=image.shape[-2:])

sample = (image, box)

translate = transforms.Lambda(functools.partial(F.affine, angle=0, translate=[-200, 0], scale=1.0, shear=0))
rotate = transforms.Lambda(functools.partial(F.rotate, angle=30, center=(0, 200)))
clamp = transforms.ClampBoundingBoxes()


vis = make_grid(
    [
        annotate(transform(sample))
        for transform in [
            lambda sample: sample,
            lambda sample: translate(sample),
            lambda sample: rotate(translate(sample)),
            lambda sample: rotate(clamp(translate(sample))),
        ]
    ],
    padding=3,
    pad_value=0,
)

However, this is somewhat unexpected. Thus, we add this clamping after all bounding box kernels that might result in objects outside the original canvas.

Running the code above with this PR gives:

Apart from the fact that the third and forth image are now equal, also note that the bounding box (red frame) is closed, whereas it is open on the left side for the image above. This shows that in the image above, the bounding box extends beyond the canvas, while it is correctly clamped here.

cc @vfdev-5 @bjuncek

[NicolasHug]

Thanks Philip, some initial feedback and questions

[vfdev-5]

In addition to affine ops, same crop would be also required by crop/pad:

import functools

import torch
from torchvision.utils import make_grid, draw_bounding_boxes
from torchvision.prototype import datapoints, transforms
from torchvision.prototype.transforms import functional as F

image = torch.full((3, 400, 400), 0.5)
image[:, 100:-100, 100:-100] = 1.0

box = datapoints.BoundingBox([100, 100, 300, 300], format="xyxy", spatial_size=image.shape[-2:])

sample = (image, box)

rotate = transforms.Lambda(functools.partial(F.rotate, angle=45, center=(200, 200)))
crop = transforms.Lambda(functools.partial(F.crop, top=200, left=200, height=400, width=400))

vis = make_grid(
    [
        annotate(transform(sample))
        for transform in [
            lambda sample: sample,
            lambda sample: crop(sample),
            lambda sample: rotate(crop(sample)),
            
        ]
    ],
    padding=3,
    pad_value=0,
)

[pmeier]

The only things left are multiple failing reference / correctness tests that currently do not have any clamping. @vfdev-5 would you happen to have some time to fix these?

[pmeier]

Thanks a ton @vfdev-5! I have a few questions inline, but this is looking almost good to go.

[vfdev-5]

We are all agree that this is good :)

[pmeier]

Although @vfdev-5 did the bulk of the work here, he needed to self-approve this PR, since I'm technically the author.