Add rotate transformation tests for rotated boxes

Test Plan:
Unit tests:
```bash
pytest test/test_transforms_v2.py -vvv -k "TestRotate and test_kernel_bounding_boxes"
pytest test/test_transforms_v2.py -vvv -k "TestRotate and test_functional_bounding_boxes_correctness"
pytest test/test_transforms_v2.py -vvv -k "TestRotate and test_transform_bounding_boxes_correctness"
```

Author: AntoineSimoulin

test/test_transforms_v2.py

@@ -1889,7 +1889,7 @@ def test_kernel_image(self, param, value, dtype, device):
         expand=[False, True],
         center=_EXHAUSTIVE_TYPE_AFFINE_KWARGS["center"],
     )
-    @pytest.mark.parametrize("format", SUPPORTED_BOX_FORMATS)
+    @pytest.mark.parametrize("format", list(tv_tensors.BoundingBoxFormat))
     @pytest.mark.parametrize("dtype", [torch.float32, torch.uint8])
     @pytest.mark.parametrize("device", cpu_and_cuda())
     def test_kernel_bounding_boxes(self, param, value, format, dtype, device):
@@ -2025,6 +2025,13 @@ def _recenter_bounding_boxes_after_expand(self, bounding_boxes, *, recenter_xy):
         x, y = recenter_xy
         if bounding_boxes.format is tv_tensors.BoundingBoxFormat.XYXY:
             translate = [x, y, x, y]
+        elif bounding_boxes.format is tv_tensors.BoundingBoxFormat.XYXYXYXY:
+            translate = [x, y, x, y, x, y, x, y]
+        elif (
+            bounding_boxes.format is tv_tensors.BoundingBoxFormat.CXCYWHR
+            or bounding_boxes.format is tv_tensors.BoundingBoxFormat.XYWHR
+        ):
+            translate = [x, y, 0.0, 0.0, 0.0]
         else:
             translate = [x, y, 0.0, 0.0]
         return tv_tensors.wrap(
@@ -2049,7 +2056,12 @@ def _reference_rotate_bounding_boxes(self, bounding_boxes, *, angle, expand, cen
             expand=expand, canvas_size=bounding_boxes.canvas_size, affine_matrix=affine_matrix
         )
 
-        output = reference_affine_bounding_boxes_helper(
+        helper = (
+            reference_affine_rotated_bounding_boxes_helper
+            if tv_tensors.is_rotated_bounding_format(bounding_boxes.format)
+            else reference_affine_bounding_boxes_helper
+        )
+        output = helper(
             bounding_boxes,
             affine_matrix=affine_matrix,
             new_canvas_size=new_canvas_size,
@@ -2060,7 +2072,7 @@ def _reference_rotate_bounding_boxes(self, bounding_boxes, *, angle, expand, cen
             bounding_boxes
         )
 
-    @pytest.mark.parametrize("format", SUPPORTED_BOX_FORMATS)
+    @pytest.mark.parametrize("format", list(tv_tensors.BoundingBoxFormat))
     @pytest.mark.parametrize("angle", _CORRECTNESS_AFFINE_KWARGS["angle"])
     @pytest.mark.parametrize("expand", [False, True])
     @pytest.mark.parametrize("center", _CORRECTNESS_AFFINE_KWARGS["center"])
@@ -2073,7 +2085,7 @@ def test_functional_bounding_boxes_correctness(self, format, angle, expand, cent
         torch.testing.assert_close(actual, expected)
         torch.testing.assert_close(F.get_size(actual), F.get_size(expected), atol=2 if expand else 0, rtol=0)
 
-    @pytest.mark.parametrize("format", SUPPORTED_BOX_FORMATS)
+    @pytest.mark.parametrize("format", list(tv_tensors.BoundingBoxFormat))
     @pytest.mark.parametrize("expand", [False, True])
     @pytest.mark.parametrize("center", _CORRECTNESS_AFFINE_KWARGS["center"])
     @pytest.mark.parametrize("seed", list(range(5)))

torchvision/transforms/v2/functional/_geometry.py

@@ -983,7 +983,7 @@ def _affine_bounding_boxes_with_expand(
         new_points = torch.matmul(points, transposed_affine_matrix)
         tr = torch.amin(new_points, dim=0, keepdim=True)
         # Translate bounding boxes
-        out_bboxes.sub_(tr.repeat((1, 2)))
+        out_bboxes.sub_(tr.repeat((1, 4 if is_rotated else 2)))
         # Estimate meta-data for image with inverted=True
         affine_vector = _get_inverse_affine_matrix(center, angle, translate, scale, shear)
         new_width, new_height = _compute_affine_output_size(affine_vector, width, height)