[proto] Added functional rotate_segmentation_mask op

[vfdev-5]

Related to #5514

Description:

• Added functional rotate_segmentation_mask op
• Added tests

Results on synthetic images/bboxes/segm mask:

Code

import numpy as np

import torch
import torchvision
from torchvision.prototype import features
from torchvision.prototype.transforms.functional import rotate_bounding_box, rotate_image_tensor, rotate_segmentation_mask

size = (64, 72)
# 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 = 22
center = None
expand = True

out_boxes = rotate_bounding_box(
    in_boxes, 
    in_boxes.format,
    in_boxes.image_size,
    angle,
    expand=expand,
    center=center
)
print(out_boxes)

out_mask = rotate_segmentation_mask(
    in_mask, 
    angle,
    expand=expand,
    center=center    
)

t_im2 = rotate_image_tensor(t_im1, angle, expand=expand, center=center)


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

plt.subplot(2,3,1)
plt.title("Input image + bboxes")
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)


plt.subplot(2,3,2)
plt.title("Input segm mask")
plt.imshow(in_mask[0, :, :].cpu().numpy())


plt.subplot(2,3,3)
plt.title("Input image + bboxes + segm mask")
plt.imshow(r1, alpha=0.5)
plt.imshow(in_mask[0, :, :].cpu().numpy(), alpha=0.75)


plt.subplot(2,3,4)
plt.title("Output image + bboxes")
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)


plt.subplot(2,3,5)
plt.title("Output segm mask")
plt.imshow(out_mask[0, :, :].cpu().numpy())

plt.subplot(2,3,6)
plt.title("Output image + bboxes + segm mask")
plt.imshow(r2, alpha=0.5)
plt.imshow(out_mask[0, :, :].cpu().numpy(), alpha=0.75)

[facebook-github-bot]

💊 CI failures summary and remediations

As of commit e0f7663 (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]

LGTM, thanks!