From 5822de4d78cd8f1a349911dfd33d751b3b692e25 Mon Sep 17 00:00:00 2001 From: Zhitao Yu Date: Mon, 23 Feb 2026 16:31:37 -0800 Subject: [PATCH 1/2] address the issue by enlarging the grid --- .../transforms/v2/functional/_geometry.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/torchvision/transforms/v2/functional/_geometry.py b/torchvision/transforms/v2/functional/_geometry.py index 84986629ba6..0e40e08653b 100644 --- a/torchvision/transforms/v2/functional/_geometry.py +++ b/torchvision/transforms/v2/functional/_geometry.py @@ -2442,10 +2442,22 @@ def elastic_bounding_boxes( convert_bounding_box_format(bounding_boxes.clone(), old_format=format, new_format=intermediate_format) ).reshape(-1, 5 if is_rotated else 4) - id_grid = _create_identity_grid(canvas_size, device=device, dtype=dtype) + # Create a grid with (H+1, W+1) to handle boundary coordinates (e.g., x=W, y=H) + extended_size = (canvas_size[0] + 1, canvas_size[1] + 1) + id_grid = _create_identity_grid(extended_size, device=device, dtype=dtype) + + # Pad displacement to match extended grid size (replicate edge values) + padded_displacement = torch.nn.functional.pad( + displacement.permute(0, 3, 1, 2), # (1, 2, H, W) + (0, 1, 0, 1), # pad right and bottom by 1 + mode="replicate", + ).permute( + 0, 2, 3, 1 + ) # back to (1, H+1, W+1, 2) + # We construct an approximation of inverse grid as inv_grid = id_grid - displacement # This is not an exact inverse of the grid - inv_grid = id_grid.sub_(displacement) + inv_grid = id_grid.sub_(padded_displacement) # Get points from bboxes points = bounding_boxes[:, :2] if is_rotated else bounding_boxes[:, [[0, 1], [2, 1], [2, 3], [0, 3]]] From 73067ede7c8c25a1d1e9d727ce1cae01ca3ea1c1 Mon Sep 17 00:00:00 2001 From: Zhitao Yu Date: Tue, 24 Feb 2026 22:34:10 -0800 Subject: [PATCH 2/2] refine the comment --- torchvision/transforms/v2/functional/_geometry.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/torchvision/transforms/v2/functional/_geometry.py b/torchvision/transforms/v2/functional/_geometry.py index 0e40e08653b..91deb67a36c 100644 --- a/torchvision/transforms/v2/functional/_geometry.py +++ b/torchvision/transforms/v2/functional/_geometry.py @@ -2448,12 +2448,12 @@ def elastic_bounding_boxes( # Pad displacement to match extended grid size (replicate edge values) padded_displacement = torch.nn.functional.pad( - displacement.permute(0, 3, 1, 2), # (1, 2, H, W) + displacement.permute(0, 3, 1, 2), # NHWC -> NCHW format (0, 1, 0, 1), # pad right and bottom by 1 mode="replicate", ).permute( 0, 2, 3, 1 - ) # back to (1, H+1, W+1, 2) + ) # back to NHWC # We construct an approximation of inverse grid as inv_grid = id_grid - displacement # This is not an exact inverse of the grid