Skip to content

Commit 550656f

Browse files
committed
cvcuda rotate verified correct visualizly and passing all tests
1 parent 8861042 commit 550656f

File tree

2 files changed

+12
-9
lines changed

2 files changed

+12
-9
lines changed

test/test_transforms_v2.py

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2204,7 +2204,8 @@ def test_functional_image_correctness(self, angle, center, interpolation, expand
22042204

22052205
actual = F.rotate(image, angle=angle, center=center, interpolation=interpolation, expand=expand, fill=fill)
22062206

2207-
if make_input == make_image_cvcuda:
2207+
if make_input is make_image_cvcuda:
2208+
actual = cvcuda_to_pil_compatible_tensor(actual)
22082209
image = cvcuda_to_pil_compatible_tensor(image)
22092210

22102211
expected = F.to_image(
@@ -2214,7 +2215,7 @@ def test_functional_image_correctness(self, angle, center, interpolation, expand
22142215
)
22152216

22162217
mae = (actual.float() - expected.float()).abs().mean()
2217-
if make_input == make_image_cvcuda:
2218+
if make_input is make_image_cvcuda:
22182219
# CV-CUDA nearest interpolation differs significantly from PIL, set much higher bound
22192220
assert mae < (122.5) if interpolation is transforms.InterpolationMode.NEAREST else 6, f"MAE: {mae}"
22202221
else:
@@ -2254,16 +2255,14 @@ def test_transform_image_correctness(self, center, interpolation, expand, fill,
22542255

22552256
torch.manual_seed(seed)
22562257

2257-
if make_input == make_image_cvcuda:
2258-
actual = F.cvcuda_to_tensor(actual).to(device="cpu")
2259-
image = F.cvcuda_to_tensor(image)
2260-
# drop the batch dimensions
2261-
image = image.squeeze(0)
2258+
if make_input is make_image_cvcuda:
2259+
actual = cvcuda_to_pil_compatible_tensor(actual)
2260+
image = cvcuda_to_pil_compatible_tensor(image)
22622261

22632262
expected = F.to_image(transform(F.to_pil_image(image)))
22642263

22652264
mae = (actual.float() - expected.float()).abs().mean()
2266-
if make_input == make_image_cvcuda:
2265+
if make_input is make_image_cvcuda:
22672266
# CV-CUDA nearest interpolation differs significantly from PIL, set much higher bound
22682267
assert mae < (122.5) if interpolation is transforms.InterpolationMode.NEAREST else 6, f"MAE: {mae}"
22692268
else:

torchvision/transforms/v2/functional/_geometry.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1596,8 +1596,11 @@ def _rotate_cvcuda(
15961596
# We need to calculate a shift to effectively rotate around the desired center
15971597
if center is None:
15981598
cx, cy = input_width / 2.0, input_height / 2.0
1599+
center_f = [0.0, 0.0]
15991600
else:
16001601
cx, cy = float(center[0]), float(center[1])
1602+
# Convert to image-center-relative coordinates (same as torchvision)
1603+
center_f = [cx - input_width * 0.5, cy - input_height * 0.5]
16011604

16021605
angle_rad = math.radians(angle)
16031606
cos_angle = math.cos(angle_rad)
@@ -1611,7 +1614,8 @@ def _rotate_cvcuda(
16111614
return cvcuda.rotate(inpt, angle_deg=angle, shift=(shift_x, shift_y), interpolation=interp)
16121615

16131616
# if we need to expand, use much of the same logic as torchvision, for output size/pad
1614-
matrix = _get_inverse_affine_matrix([0.0, 0.0], -angle, [0.0, 0.0], 1.0, [0.0, 0.0])
1617+
# Use center_f (image-center-relative coords) to match torchvision's output size calculation
1618+
matrix = _get_inverse_affine_matrix(center_f, -angle, [0.0, 0.0], 1.0, [0.0, 0.0])
16151619
output_width, output_height = _compute_affine_output_size(matrix, input_width, input_height)
16161620

16171621
pad_left = (output_width - input_width) // 2

0 commit comments

Comments
 (0)