Skip to content

Commit d5df0d6

Browse files
Rotated bboxes transforms (#9104)
1 parent 6473b77 commit d5df0d6

File tree

6 files changed

+410
-131
lines changed

6 files changed

+410
-131
lines changed

test/common_utils.py

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121
from torch.testing._comparison import BooleanPair, NonePair, not_close_error_metas, NumberPair, TensorLikePair
2222
from torchvision import io, tv_tensors
2323
from torchvision.transforms._functional_tensor import _max_value as get_max_value
24-
from torchvision.transforms.v2.functional import to_image, to_pil_image
24+
from torchvision.transforms.v2.functional import clamp_bounding_boxes, to_image, to_pil_image
2525

2626

2727
IN_OSS_CI = any(os.getenv(var) == "true" for var in ["CIRCLECI", "GITHUB_ACTIONS"])
@@ -467,9 +467,20 @@ def sample_position(values, max_value):
467467
parts = (x1, y1, x2, y2, x3, y3, x4, y4)
468468
else:
469469
raise ValueError(f"Format {format} is not supported")
470-
return tv_tensors.BoundingBoxes(
471-
torch.stack(parts, dim=-1).to(dtype=dtype, device=device), format=format, canvas_size=canvas_size
472-
)
470+
out_boxes = torch.stack(parts, dim=-1).to(dtype=dtype, device=device)
471+
if tv_tensors.is_rotated_bounding_format(format):
472+
# The rotated bounding boxes are not guaranteed to be within the canvas by design,
473+
# so we apply clamping. We also add a 2 buffer to the canvas size to avoid
474+
# numerical issues during the testing
475+
buffer = 4
476+
out_boxes = clamp_bounding_boxes(
477+
out_boxes, format=format, canvas_size=(canvas_size[0] - buffer, canvas_size[1] - buffer)
478+
)
479+
if format is tv_tensors.BoundingBoxFormat.XYWHR or format is tv_tensors.BoundingBoxFormat.CXCYWHR:
480+
out_boxes[:, :2] += buffer // 2
481+
elif format is tv_tensors.BoundingBoxFormat.XYXYXYXY:
482+
out_boxes[:, :] += buffer // 2
483+
return tv_tensors.BoundingBoxes(out_boxes, format=format, canvas_size=canvas_size)
473484

474485

475486
def make_detection_masks(size=DEFAULT_SIZE, *, num_masks=1, dtype=None, device="cpu"):

0 commit comments

Comments
 (0)