[proto] Added functional affine_bounding_box op

[vfdev-5]

Related to #5514

Description:

• Added functional affine_bounding_box op
• Added tests
  • 579 tests and can take 5.83s (on my machine)

Questions:

1. For actual code and the stable API, positive angle parameter (affine op) performs clockwise rotation. For D2 and albu does counter clockwise rotation. Should we align our proto API to that ? Pillow rotate and the stable API rotate op does also counter clockwise rotation.

vision/torchvision/transforms/functional.py

Line 1004 in 8aadef5

angle (number): rotation angle value in degrees, counter-clockwise.

vision/torchvision/transforms/functional.py

Line 1092 in 8aadef5

angle (number): rotation angle in degrees between -180 and 180, clockwise direction.

=>
We assume that affine_image_tensor == F.affine and must be BC, so let's keep current behaviour.
Opened related RFC: #5608

Compare with albumentations:

Code and results. TL;DR: results should match

# pip install albumentations

import numpy as np
from albumentations.augmentations.geometric.functional import bbox_shift_scale_rotate, normalize_bbox, denormalize_bbox

import torch
import torchvision
from torchvision.prototype import features
from torchvision.prototype.transforms.functional import affine_bounding_box

print(torch.__version__)
print(torchvision.__version__)


size = (64, 76)
# xyxy format
in_boxes = [
    [20, 25, 35, 45],
    [50, 5, 70, 22],
]

im1 = 255 * np.ones(size + (3, ), dtype=np.uint8)
for in_box in in_boxes:
    im1[in_box[1]:in_box[3], in_box[0]:in_box[2], :] = (127, 127, 127)


# Params
angle = 63
scale = 0.89
dx = 0.12
dy = 0.23


print("")
print("-- albu --")

albu_out_boxes = []
for in_box in in_boxes:
    n_in_box = normalize_bbox(in_box, *size)
    n_out_box = bbox_shift_scale_rotate(n_in_box, -angle, scale, dx, dy, *size)
    out_box = denormalize_bbox(n_out_box, *size)
    albu_out_boxes.append(out_box)

print("-- output:")
for out_box in albu_out_boxes:
    print("\t", out_box)


print("")
print("-- vision --")

in_boxes = features.BoundingBox(
    in_boxes, format=features.BoundingBoxFormat.XYXY, image_size=size, dtype=torch.float64
)

out_boxes = affine_bounding_box(
    in_boxes,
    in_boxes.format,
    in_boxes.image_size,
    angle,
    (dx * size[1], dy * size[0]),
    scale,
    shear=(0, 0)
)

print("-- output:")
for out_box in out_boxes:
    print("\t", out_box.cpu().tolist())


for a_out_box, out_box in zip(albu_out_boxes, out_boxes):
    np.testing.assert_allclose(
        out_box.cpu().numpy(), a_out_box
    )

Output:

1.12.0.dev20220311+cu113
0.13.0a0+833fa10

-- albu --
-- output:
         (29.538126709313037, 29.61771466912497, 51.45881601138894, 49.593682662403616)
         (59.89857660249483, 45.32655796959046, 81.46053620882876, 68.05535036120276)

-- vision --
-- output:
         [29.53812670931304, 29.61771466912497, 51.45881601138894, 49.593682662403616]
         [59.898576602494835, 45.32655796959046, 81.46053620882877, 68.05535036120274]

Results on synthetic images/bboxes:

Code

import numpy as np
from albumentations.augmentations.geometric.functional import bbox_shift_scale_rotate, normalize_bbox, denormalize_bbox
import cv2
import matplotlib.pyplot as plt

import torch
import torchvision
from torchvision.prototype import features
from torchvision.prototype.transforms.functional import affine_bounding_box, affine_image_tensor


size = (64, 76)
# xyxy format
in_boxes = [
    [10, 15, 25, 35],
    [50, 5, 70, 22],
    [45, 46, 56, 62],
]

im1 = 255 * np.ones(size + (3, ), dtype=np.uint8)
for in_box in in_boxes:
    im1[in_box[1]:in_box[3], in_box[0]:in_box[2], :] = (127, 127, 127)

t_im1 = torch.tensor(im1).permute(2, 0, 1).view(1, 3, *size)

in_boxes = features.BoundingBox(
    in_boxes, format=features.BoundingBoxFormat.XYXY, image_size=size
)
    
angle = 34
scale = 0.9
t = (-6, 7)
shear = (1, 2)

out_boxes = affine_bounding_box(
    in_boxes, 
    in_boxes.format,
    in_boxes.image_size,
    angle,
    t,
    scale,
    shear,
)

t_im2 = affine_image_tensor(t_im1, angle, t, scale, shear)scale = 0.9

plt.figure(figsize=(14, 8))

ax = plt.subplot(1,2,1)
r1 = t_im1[0, ...].permute(1, 2, 0).contiguous().cpu().numpy()
for in_box in in_boxes:    
    r1 = cv2.rectangle(r1, (in_box[0].item(), in_box[1].item()), (in_box[2].item(), in_box[3].item()), (255, 127, 0))
plt.imshow(r1)


ax = plt.subplot(1,2,2)
r2 = t_im2[0, ...].permute(1, 2, 0).contiguous().cpu().numpy()
for out_box in out_boxes:
    out_box = np.round(out_box.cpu().numpy()).astype("int32")
    r2 = cv2.rectangle(r2, (out_box[0], out_box[1]), (out_box[2], out_box[3]), (255, 127, 0), 0)
plt.imshow(r2)

Output boxes

out_boxes:
tensor([[ 9.0663, 13.0952, 30.8541, 35.0096],
        [46.3305, 24.8059, 70.3870, 46.8929],
        [21.8478, 52.6698, 38.5143, 69.7240]])

[facebook-github-bot]

💊 CI failures summary and remediations

As of commit c398dad (more details on the Dr. CI page):

---

💚 💚 Looks good so far! There are no failures yet. 💚 💚

---

This comment was automatically generated by Dr. CI (expand for details).

Please report bugs/suggestions to the (internal) Dr. CI Users group.

Click here to manually regenerate this comment.

[datumbox]

Thanks @vfdev-5. Few very minor comments and a couple of questions. Let me know what you think.

[datumbox]

LGTM, thanks!

[vfdev-5]

Refactored this method as it will be reused for affine segm mask tests

[datumbox]

LGTM, thanks!