[proto] Added functional crop_bounding_box op

[vfdev-5]

Related to #5514

Description:

• Added functional crop_bounding_box op
• Added tests

Results match albumentations' crop_bbox_by_coords

Results on synthetic images/bboxes:

Code

import numpy as np

import torch
import torchvision
from torchvision.prototype import features
from torchvision.prototype.transforms.functional import crop_bounding_box, crop_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 i, in_box in enumerate(in_boxes):
    im1[in_box[1]:in_box[3], in_box[0]:in_box[2], :] = (i * 40, 200 - i * 40, 127 + i * 40)

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
)
    
t, l, h, w = -8, 12, 80, 40


out_boxes = crop_bounding_box(
    in_boxes, 
    in_boxes.format,
    in_boxes.image_size,
    t,
    l,
    h,
    w,
)

t_im2 = crop_image_tensor(t_im1, t, l, h, w).clone()

import cv2
import matplotlib.pyplot as plt
%matplotlib inline


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)

[datumbox]

Thanks @vfdev-5, just one question.

@pmeier Thoughts on the testing strategy?

[pmeier]

One comment inline. Testing looks ok.

[datumbox]

@vfdev-5 Thanks a lot for the contribution. LGTM.

@pmeier any concerns from you or do we merge?

[pmeier]

Thanks Victor!

[github-actions]

Hey @pmeier!

You merged this PR, but no labels were added. The list of valid labels is available at https://github.com/pytorch/vision/blob/main/.github/process_commit.py

[vadimkantorov]

For batch support it may be better to instead use bounding_box[..., 0::2]. Then view(-1, 4) and view(shape) would not be needed?

[pmeier]

Good catch. I was under the impression ... was not supported by torch.jit.script in general, but this seems fine:

@torch.jit.script
def foo(x: torch.Tensor) -> torch.Tensor:
    y = x.clone()
    y[..., 0::2] += 1
    return y

It seems, only the explicit indices are not supported:

@torch.jit.script
def bar(x: torch.Tensor) -> torch.Tensor:
    y = x.clone()
    y[..., [0, 2]] += 1
    return y

RuntimeError: 
Ellipses followed by tensor indexing is currently not supported:
[...]
def bar(x: torch.Tensor) -> torch.Tensor:
    y = x.clone()
    y[..., [0, 2]] += 1
    ~~~~~~~~~~~~~ <--- HERE
    return y