-
Notifications
You must be signed in to change notification settings - Fork 7.2k
Add 'crop' option to RandomRotation #9458
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -2462,6 +2462,62 @@ def test_functional_image_fast_path_correctness(self, size, angle, expand): | |
|
|
||
| torch.testing.assert_close(actual, expected) | ||
|
|
||
| @pytest.mark.parametrize("size", [(100, 100), (120, 80)]) | ||
| @pytest.mark.parametrize("angle", [15.0, 30.0, 45.0]) | ||
| def test_transform_crop_removes_fill(self, size, angle): | ||
| # Output of crop=True should contain no fill pixels when input is fully non-zero | ||
| h, w = size | ||
| image = tv_tensors.Image(torch.full((3, h, w), 200, dtype=torch.uint8)) | ||
| transform = transforms.RandomRotation((angle, angle), fill=0, crop=True) | ||
| output = transform(image) | ||
| assert output.min().item() > 0, "crop=True output should have no fill pixels" | ||
| assert output.shape[-2] < h or output.shape[-1] < w, "crop=True should reduce at least one dimension" | ||
|
|
||
| @pytest.mark.parametrize("size", [(100, 100), (120, 80)]) | ||
| @pytest.mark.parametrize("angle", [15.0, 30.0, 45.0]) | ||
| def test_transform_crop_consistent_across_inputs(self, size, angle): | ||
| # Image, mask, and bounding boxes should all be cropped to the same canvas size | ||
| h, w = size | ||
| image = tv_tensors.Image(torch.full((3, h, w), 200, dtype=torch.uint8)) | ||
| mask = tv_tensors.Mask(torch.ones(1, h, w, dtype=torch.uint8)) | ||
| boxes = tv_tensors.BoundingBoxes( | ||
| torch.tensor([[10.0, 10.0, 50.0, 50.0]]), | ||
| format=tv_tensors.BoundingBoxFormat.XYXY, | ||
| canvas_size=(h, w), | ||
| ) | ||
| transform = transforms.RandomRotation((angle, angle), crop=True) | ||
| out_image, out_mask, out_boxes = transform(image, mask, boxes) | ||
| assert out_image.shape[-2:] == out_mask.shape[-2:] | ||
| assert out_boxes.canvas_size == (out_image.shape[-2], out_image.shape[-1]) | ||
|
|
||
| def test_transform_crop_and_expand_mutually_exclusive(self): | ||
| with pytest.raises(ValueError, match="crop and expand are mutually exclusive"): | ||
| transforms.RandomRotation(30, expand=True, crop=True) | ||
|
|
||
| @pytest.mark.parametrize("angle", [0.0, 90.0, 180.0, 270.0]) | ||
| def test_transform_crop_zero_angle_preserves_size(self, angle): | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The test name "test_transform_crop_zero_angle_preserves_size" says "zero_angle" but it tests 0°, 90°, 180°, and 270°, which is a bit misleading. Could you please rename it?
Comment on lines
+2497
to
+2498
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Can we add non-square sizes by parametrizing over non-square sizes like (120, 80) here? |
||
| # Multiples of 90° should not reduce the image size | ||
| image = tv_tensors.Image(torch.zeros(3, 100, 100, dtype=torch.uint8)) | ||
| transform = transforms.RandomRotation((angle, angle), crop=True) | ||
| output = transform(image) | ||
| assert output.shape == image.shape | ||
|
|
||
| def test_largest_inscribed_crop_size(self): | ||
| from torchvision.transforms.v2.functional._geometry import _largest_inscribed_crop_size | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. We import the functions on the top of the file. For consistency, please move the |
||
|
|
||
| # No rotation: crop equals original size | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This comment is obvious from the code, so we can remove it. |
||
| assert _largest_inscribed_crop_size(100, 100, 0) == (100, 100) | ||
| assert _largest_inscribed_crop_size(200, 100, 0) == (100, 200) | ||
|
|
||
| # 45° square: inscribed square has side = 100 / sqrt(2) ≈ 70.71 → floor to 70 | ||
| crop_h, crop_w = _largest_inscribed_crop_size(100, 100, 45) | ||
| assert crop_h == crop_w == 70 | ||
|
|
||
| # Crop is always smaller than or equal to original dimensions | ||
| for w, h, a in [(200, 100, 20), (640, 480, 15), (50, 50, 37)]: | ||
| ch, cw = _largest_inscribed_crop_size(w, h, a) | ||
| assert ch <= h and cw <= w | ||
|
|
||
|
|
||
| class TestContainerTransforms: | ||
| class BuiltinTransform(transforms.Transform): | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -1343,6 +1343,48 @@ def affine_video( | |
| ) | ||
|
|
||
|
|
||
| def _largest_inscribed_crop_size(width: int, height: int, angle: float) -> tuple[int, int]: | ||
| """Compute the largest axis-aligned rectangle inscribed in a rotated width x height rectangle. | ||
|
|
||
| Returns ``(crop_height, crop_width)`` as integers. | ||
|
|
||
| Approach taken from https://stackoverflow.com/questions/16702966/rotate-image-and-crop-out-black-borders | ||
| """ | ||
| angle_rad = math.radians(angle) | ||
| sin_a = abs(math.sin(angle_rad)) | ||
| cos_a = abs(math.cos(angle_rad)) | ||
|
|
||
| # Guard against small values of sin_a or cos_a that will cause us to truncate a pixel | ||
| # off rotations that introduce no padding (e.g. 90° or 180°). | ||
| if sin_a < 1e-10: | ||
| return height, width | ||
| if cos_a < 1e-10: | ||
| return width, height | ||
|
|
||
| width_is_longer = width >= height | ||
| side_long = width if width_is_longer else height | ||
| side_short = height if width_is_longer else width | ||
|
|
||
| if side_short <= 2.0 * sin_a * cos_a * side_long or abs(sin_a - cos_a) < 1e-10: | ||
| # Half-constrained: two crop corners touch the longer side. | ||
| # Also handles the 45° case where sin_a == cos_a and the denominator | ||
| # in the fully constrained solution goes to zero. | ||
| x = 0.5 * side_short | ||
| if width_is_longer: | ||
| crop_w, crop_h = x / sin_a, x / cos_a | ||
| else: | ||
| crop_w, crop_h = x / cos_a, x / sin_a | ||
| else: | ||
| # Fully constrained: crop touches all four sides | ||
| cos_2a = cos_a * cos_a - sin_a * sin_a | ||
| crop_w = (width * cos_a - height * sin_a) / cos_2a | ||
| crop_h = (height * cos_a - width * sin_a) / cos_2a | ||
|
|
||
| # Use floor (int()) to guarantee the crop region contains no fill pixels. | ||
| # Clamp to image dimensions for edge cases like wide images rotated near 90°. | ||
| return min(int(crop_h), height), min(int(crop_w), width) | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Would math.floor() be more appropriate here? int() truncates toward zero rather than rounding down, and math.floor() would make the "round down" intent more explicit. (Both behave the same for positive values, which is always the case here, so this is just a readability suggestion.) |
||
|
|
||
|
|
||
| def rotate( | ||
| inpt: torch.Tensor, | ||
| angle: float, | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We don't need specify anything beyond the assert, pytest is already good at showing the right thing. Please remove the message after assert.