Added functional affine_bounding_box op with tests

Author: vfdev-5

test/test_prototype_transforms_functional.py

@@ -1,14 +1,18 @@
 import functools
 import itertools
+import math
 
+import numpy as np
 import pytest
 import torch.testing
 import torchvision.prototype.transforms.functional as F
 from torch import jit
 from torch.nn.functional import one_hot
 from torchvision.prototype import features
+from torchvision.prototype.transforms.functional._meta import convert_bounding_box_format
 from torchvision.transforms.functional_tensor import _max_value as get_max_value
 
+
 make_tensor = functools.partial(torch.testing.make_tensor, device="cpu")
 
 
@@ -205,6 +209,45 @@ def resize_bounding_box():
             yield SampleInput(bounding_box, size=size, image_size=bounding_box.image_size)
 
 
+@register_kernel_info_from_sample_inputs_fn
+def affine_image_tensor():
+    for image, angle, translate, scale, shear in itertools.product(
+        make_images(extra_dims=()),
+        [-87, 15, 90],  # angle
+        [5, -5],  # translate
+        [0.77, 1.27],  # scale
+        [0, 12],  # shear
+    ):
+        yield SampleInput(
+            image,
+            angle=angle,
+            translate=(translate, translate),
+            scale=scale,
+            shear=(shear, shear),
+            interpolation=F.InterpolationMode.NEAREST,
+        )
+
+
+@register_kernel_info_from_sample_inputs_fn
+def affine_bounding_box():
+    for bounding_box, angle, translate, scale, shear in itertools.product(
+        make_bounding_boxes(),
+        [-87, 15, 90],  # angle
+        [5, -5],  # translate
+        [0.77, 1.27],  # scale
+        [0, 12],  # shear
+    ):
+        yield SampleInput(
+            bounding_box,
+            format=bounding_box.format,
+            image_size=bounding_box.image_size,
+            angle=angle,
+            translate=(translate, translate),
+            scale=scale,
+            shear=(shear, shear),
+        )
+
+
 @pytest.mark.parametrize(
     "kernel",
     [
@@ -233,3 +276,98 @@ def test_eager_vs_scripted(functional_info, sample_input):
     scripted = jit.script(functional_info.functional)(*sample_input.args, **sample_input.kwargs)
 
     torch.testing.assert_close(eager, scripted)
+
+
+@pytest.mark.parametrize("angle", range(-90, 90, 36))
+@pytest.mark.parametrize("translate", range(-10, 10, 5))
+@pytest.mark.parametrize("scale", [0.77, 1.0, 1.27])
+@pytest.mark.parametrize("shear", range(-15, 15, 5))
+@pytest.mark.parametrize("center", [None, (12, 14)])
+def test_correctness_affine_bounding_box(angle, translate, scale, shear, center):
+    def _compute_expected_bbox(bbox, angle_, translate_, scale_, shear_, center_):
+        rot = math.radians(angle_)
+        cx, cy = center_
+        tx, ty = translate_
+        sx, sy = [math.radians(sh_) for sh_ in shear_]
+
+        c_matrix = np.array([[1, 0, cx], [0, 1, cy], [0, 0, 1]])
+        t_matrix = np.array([[1, 0, tx], [0, 1, ty], [0, 0, 1]])
+        c_matrix_inv = np.linalg.inv(c_matrix)
+        rs_matrix = np.array(
+            [
+                [scale_ * math.cos(rot), -scale_ * math.sin(rot), 0],
+                [scale_ * math.sin(rot), scale_ * math.cos(rot), 0],
+                [0, 0, 1],
+            ]
+        )
+        shear_x_matrix = np.array([[1, -math.tan(sx), 0], [0, 1, 0], [0, 0, 1]])
+        shear_y_matrix = np.array([[1, 0, 0], [-math.tan(sy), 1, 0], [0, 0, 1]])
+        rss_matrix = np.matmul(rs_matrix, np.matmul(shear_y_matrix, shear_x_matrix))
+        true_matrix = np.matmul(t_matrix, np.matmul(c_matrix, np.matmul(rss_matrix, c_matrix_inv)))
+        true_matrix = true_matrix[:2, :]
+
+        bbox_xyxy = convert_bounding_box_format(
+            bbox, old_format=bbox.format, new_format=features.BoundingBoxFormat.XYXY
+        )
+        points = np.array(
+            [
+                [bbox_xyxy[0].item(), bbox_xyxy[1].item(), 1.0],
+                [bbox_xyxy[2].item(), bbox_xyxy[1].item(), 1.0],
+                [bbox_xyxy[0].item(), bbox_xyxy[3].item(), 1.0],
+                [bbox_xyxy[2].item(), bbox_xyxy[3].item(), 1.0],
+            ]
+        )
+        transformed_points = points @ true_matrix.T
+        out_bbox = [
+            np.min(transformed_points[:, 0]),
+            np.min(transformed_points[:, 1]),
+            np.max(transformed_points[:, 0]),
+            np.max(transformed_points[:, 1]),
+        ]
+        out_bbox = features.BoundingBox(
+            out_bbox, format=features.BoundingBoxFormat.XYXY, image_size=(32, 32), dtype=torch.float32
+        )
+        out_bbox = convert_bounding_box_format(
+            out_bbox, old_format=features.BoundingBoxFormat.XYXY, new_format=bbox.format
+        )
+        return out_bbox
+
+    image_size = (32, 32)
+
+    for bboxes in make_bounding_boxes(
+        image_sizes=[
+            image_size,
+        ],
+        extra_dims=((4,),),
+    ):
+        output_bboxes = F.affine_bounding_box(
+            bboxes,
+            bboxes.format,
+            image_size=image_size,
+            angle=angle,
+            translate=(translate, translate),
+            scale=scale,
+            shear=(shear, shear),
+            center=center,
+        )
+        if center is None:
+            center = [s // 2 for s in image_size]
+
+        bboxes_format = bboxes.format
+        bboxes_image_size = bboxes.image_size
+        if bboxes.ndim < 2:
+            bboxes = [
+                bboxes,
+            ]
+
+        expected_bboxes = []
+        for bbox in bboxes:
+            bbox = features.BoundingBox(bbox, format=bboxes_format, image_size=bboxes_image_size)
+            expected_bboxes.append(
+                _compute_expected_bbox(bbox, angle, (translate, translate), scale, (shear, shear), center)
+            )
+        expected_bboxes = torch.stack(expected_bboxes)
+        if expected_bboxes.shape[0] < 2:
+            expected_bboxes = expected_bboxes.squeeze(0)
+
+        torch.testing.assert_close(output_bboxes, expected_bboxes)

torchvision/prototype/transforms/functional/__init__.py

@@ -48,6 +48,7 @@
     center_crop_image_pil,
     resized_crop_image_tensor,
     resized_crop_image_pil,
+    affine_bounding_box,
     affine_image_tensor,
     affine_image_pil,
     rotate_image_tensor,

torchvision/prototype/transforms/functional/_geometry.py

@@ -174,6 +174,53 @@ def affine_image_pil(
     return _FP.affine(img, matrix, interpolation=pil_modes_mapping[interpolation], fill=fill)
 
 
+def affine_bounding_box(
+    bounding_box: torch.Tensor,
+    format: features.BoundingBoxFormat,
+    image_size: Tuple[int, int],
+    angle: float,
+    translate: List[float],
+    scale: float,
+    shear: List[float],
+    center: Optional[List[float]] = None,
+) -> torch.Tensor:
+    original_shape = bounding_box.shape
+    bounding_box = convert_bounding_box_format(
+        bounding_box, old_format=format, new_format=features.BoundingBoxFormat.XYXY
+    ).view(-1, 4)
+
+    dtype = bounding_box.dtype if torch.is_floating_point(bounding_box) else torch.float32
+    device = bounding_box.device
+
+    if center is None:
+        height, width = image_size
+        center_f = [width * 0.5, height * 0.5]
+    else:
+        center_f = [float(c) for c in center]
+
+    translate_f = [float(t) for t in translate]
+    affine_matrix = torch.tensor(
+        _get_inverse_affine_matrix(center_f, angle, translate_f, scale, shear, inverted=False),
+        dtype=dtype,
+        device=device,
+    ).view(2, 3)
+    # bboxes to 4 points like:
+    # [(xmin, ymin, 1), (xmax, ymin, 1), (xmax, ymax, 1), (xmin, ymax, 1), ...]
+    points = bounding_box[:, [[0, 1], [2, 1], [2, 3], [0, 3]]].view(-1, 2)
+    points = torch.cat([points, torch.ones(points.shape[0], 1)], dim=-1)
+    transformed_points = points @ affine_matrix.T
+    # reshape transformed points to [N boxes, 4 points, x/y coords]
+    transformed_points = transformed_points.view(-1, 4, 2)
+    # compute bounding box from 4 transformed points:
+    out_bbox_mins, _ = torch.min(transformed_points, dim=1)
+    out_bbox_maxs, _ = torch.max(transformed_points, dim=1)
+    out_bboxes = torch.cat([out_bbox_mins, out_bbox_maxs], dim=1)
+    # out_bboxes should be of shape [N boxes, 4]
+    return convert_bounding_box_format(out_bboxes, old_format=features.BoundingBoxFormat.XYXY, new_format=format).view(
+        original_shape
+    )
+
+
 def rotate_image_tensor(
     img: torch.Tensor,
     angle: float,

torchvision/transforms/functional.py

@@ -931,11 +931,7 @@ def adjust_gamma(img: Tensor, gamma: float, gain: float = 1) -> Tensor:
 
 
 def _get_inverse_affine_matrix(
-    center: List[float],
-    angle: float,
-    translate: List[float],
-    scale: float,
-    shear: List[float],
+    center: List[float], angle: float, translate: List[float], scale: float, shear: List[float], inverted: bool = True
 ) -> List[float]:
     # Helper method to compute inverse matrix for affine transformation
 
@@ -970,18 +966,26 @@ def _get_inverse_affine_matrix(
     c = math.sin(rot - sy) / math.cos(sy)
     d = -math.sin(rot - sy) * math.tan(sx) / math.cos(sy) + math.cos(rot)
 
-    # Inverted rotation matrix with scale and shear
-    # det([[a, b], [c, d]]) == 1, since det(rotation) = 1 and det(shear) = 1
-    matrix = [d, -b, 0.0, -c, a, 0.0]
-    matrix = [x / scale for x in matrix]
-
-    # Apply inverse of translation and of center translation: RSS^-1 * C^-1 * T^-1
-    matrix[2] += matrix[0] * (-cx - tx) + matrix[1] * (-cy - ty)
-    matrix[5] += matrix[3] * (-cx - tx) + matrix[4] * (-cy - ty)
-
-    # Apply center translation: C * RSS^-1 * C^-1 * T^-1
-    matrix[2] += cx
-    matrix[5] += cy
+    if inverted:
+        # Inverted rotation matrix with scale and shear
+        # det([[a, b], [c, d]]) == 1, since det(rotation) = 1 and det(shear) = 1
+        matrix = [d, -b, 0.0, -c, a, 0.0]
+        matrix = [x / scale for x in matrix]
+        # Apply inverse of translation and of center translation: RSS^-1 * C^-1 * T^-1
+        matrix[2] += matrix[0] * (-cx - tx) + matrix[1] * (-cy - ty)
+        matrix[5] += matrix[3] * (-cx - tx) + matrix[4] * (-cy - ty)
+        # Apply center translation: C * RSS^-1 * C^-1 * T^-1
+        matrix[2] += cx
+        matrix[5] += cy
+    else:
+        matrix = [a, b, 0.0, c, d, 0.0]
+        matrix = [x * scale for x in matrix]
+        # Apply inverse of center translation: RSS * C^-1
+        matrix[2] += matrix[0] * (-cx) + matrix[1] * (-cy)
+        matrix[5] += matrix[3] * (-cx) + matrix[4] * (-cy)
+        # Apply translation and center : T * C * RSS * C^-1
+        matrix[2] += cx + tx
+        matrix[5] += cy + ty
 
     return matrix