From b2793619222c0c44bd4801f51be1fad7b3d96b0c Mon Sep 17 00:00:00 2001 From: Philip Meier Date: Fri, 22 Apr 2022 08:18:35 +0200 Subject: [PATCH] simplify boundingbox kernels with jit compatible ellipsis --- .../transforms/functional/_geometry.py | 20 ++++++++----------- 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/torchvision/prototype/transforms/functional/_geometry.py b/torchvision/prototype/transforms/functional/_geometry.py index 71be0a22c00..58630cdaa9b 100644 --- a/torchvision/prototype/transforms/functional/_geometry.py +++ b/torchvision/prototype/transforms/functional/_geometry.py @@ -401,18 +401,16 @@ def pad_bounding_box( ) -> torch.Tensor: left, _, top, _ = _FT._parse_pad_padding(padding) - shape = bounding_box.shape - bounding_box = convert_bounding_box_format( bounding_box, old_format=format, new_format=features.BoundingBoxFormat.XYXY - ).view(-1, 4) + ) - bounding_box[:, 0::2] += left - bounding_box[:, 1::2] += top + bounding_box[..., 0::2] += left + bounding_box[..., 1::2] += top return convert_bounding_box_format( bounding_box, old_format=features.BoundingBoxFormat.XYXY, new_format=format, copy=False - ).view(shape) + ) crop_image_tensor = _FT.crop @@ -425,19 +423,17 @@ def crop_bounding_box( top: int, left: int, ) -> torch.Tensor: - shape = bounding_box.shape - bounding_box = convert_bounding_box_format( bounding_box, old_format=format, new_format=features.BoundingBoxFormat.XYXY - ).view(-1, 4) + ) # Crop or implicit pad if left and/or top have negative values: - bounding_box[:, 0::2] -= left - bounding_box[:, 1::2] -= top + bounding_box[..., 0::2] -= left + bounding_box[..., 1::2] -= top return convert_bounding_box_format( bounding_box, old_format=features.BoundingBoxFormat.XYXY, new_format=format, copy=False - ).view(shape) + ) def perspective_image_tensor(