[proto] Added functional affine_segmentation_mask op

[vfdev-5]

Related to #5514

Description:

• Added functional affine_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 affine_image_tensor, affine_bounding_box, affine_segmentation_mask

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

im1 = 255 * np.ones(size + (3, ), dtype=np.uint8)
mask = np.zeros(size, dtype=np.int64)
for in_box, label in zip(in_boxes, labels):
    im1[in_box[1]:in_box[3], in_box[0]:in_box[2], :] = (127, 127, 127)
    mask[in_box[1]:in_box[3], in_box[0]:in_box[2]] = label
    
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
)
in_mask = features.SegmentationMask(torch.tensor(mask)).view(1, *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,
)
print(out_boxes)

out_mask = affine_segmentation_mask(
    in_mask, 
    angle,
    t,
    scale,
    shear    
)

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


import cv2
import matplotlib.pyplot as plt
%matplotlib inline


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)

Compare to albumentations:

TL;DR: results do not match due to missing offset and opencv/opencv#11784

Code

# pip install albumentations
# pip install git+https://github.com/pytorch/data

import numpy as np
import cv2
import albumentations
from albumentations.augmentations.geometric.functional import shift_scale_rotate

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

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


size = (64, 64)
# xyxy format
in_boxes = [
    [50, 5, 70, 22],
    [size[1] // 2 - 10, size[0] // 2 - 10, size[1] // 2 + 10, size[0] // 2 + 10],
    [1, 1, 5, 5],
]
labels = [1, 2, 3]

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


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


# https://github.com/albumentations-team/albumentations/blob/89a675cbfb2b76f6be90e7049cd5211cb08169a5/albumentations/augmentations/geometric/transforms.py#L81
albu_out_mask = shift_scale_rotate(mask, -angle, scale, dx, dy, cv2.INTER_NEAREST, cv2.BORDER_CONSTANT, 0)

# Using offset for images
# https://github.com/opencv/opencv/issues/11784
offset = 0.5
center = (size[1] / 2 - offset, size[0] / 2 - offset)
m = cv2.getRotationMatrix2D(center, -angle, scale=scale)
m[0, 2] += dx * size[1]
m[1, 2] += dy * size[0]
cv2_out_mask = cv2.warpAffine(mask, m, dsize=size[::-1], flags=cv2.INTER_NEAREST, borderValue=0, borderMode=0)

in_mask = features.SegmentationMask(torch.tensor(mask)).view(1, *size)

out_mask = affine_segmentation_mask(
    in_mask,
    angle,
    (dx * size[1], dy * size[0]),
    scale,
    shear=(0, 0)
)


import matplotlib.pyplot as plt
%matplotlib inline

plt.figure(figsize=(20, 15))
plt.subplot(141)
plt.title("Input mask")
plt.imshow(mask)
plt.subplot(142)
plt.title("Output mask by torchvision")
plt.imshow(np_out_mask)
plt.subplot(143)
plt.title("Output mask by albumentations")
plt.imshow(albu_out_mask)
plt.subplot(144)
plt.title("Output mask diff: torchvision - albumentations")
plt.imshow(np_out_mask - albu_out_mask)


plt.figure(figsize=(20, 15))
plt.subplot(141)
plt.title("Input mask")
plt.imshow(mask)
plt.subplot(142)
plt.title("Output mask by torchvision")
plt.imshow(np_out_mask)
plt.subplot(143)
plt.title("Output mask by offsetted cv2 affine warp")
plt.imshow(cv2_out_mask)
plt.subplot(144)
plt.title("Output mask diff: torchvision - offsetted cv2 affine warp")
plt.imshow(np_out_mask - cv2_out_mask)

[facebook-github-bot]

💊 CI failures summary and remediations

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

---

✅ None of the CI failures appear to be your fault 💚

---

• 3/3 broken upstream at merge base 65d3a87 since Mar 21

---

🚧 3 ongoing upstream failures:

These were probably caused by upstream breakages that are not fixed yet.

• binary_linux_conda_py3.8_cu115 since Mar 21 (fbc8ea4)
  • 🔁 rerun
• binary_linux_conda_py3.7_cu115 since Mar 21 (fbc8ea4)
  • 🔁 rerun
• binary_linux_conda_py3.10_cu115 since Mar 21 (fbc8ea4)
  • 🔁 rerun

---

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]

@vfdev-5 LGTM on the kernel side.

@pmeier Are your concerns on the test side covered or you would suggest more changes?

[pmeier]

I've added some more test related comments.

[pmeier]

Thanks @vfdev-5! LGTM when CI is green.

[github-actions]

Hey @vfdev-5!

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