From dbf080a051d4f6d4203ba702402a73b063c93f16 Mon Sep 17 00:00:00 2001 From: Philip Meier Date: Wed, 26 Oct 2022 08:44:52 +0200 Subject: [PATCH] reinstate _split_alpha --- torchvision/prototype/transforms/functional/_meta.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/torchvision/prototype/transforms/functional/_meta.py b/torchvision/prototype/transforms/functional/_meta.py index 5e017848415..bfa7d8bc9e9 100644 --- a/torchvision/prototype/transforms/functional/_meta.py +++ b/torchvision/prototype/transforms/functional/_meta.py @@ -183,6 +183,10 @@ def clamp_bounding_box( return convert_format_bounding_box(xyxy_boxes, BoundingBoxFormat.XYXY, format) +def _split_alpha(image: torch.Tensor) -> List[torch.Tensor]: + return torch.tensor_split(image, indices=(-1,), dim=-3) + + def _strip_alpha(image: torch.Tensor) -> torch.Tensor: image, alpha = torch.tensor_split(image, indices=(-1,), dim=-3) if not torch.all(alpha == _FT._max_value(alpha.dtype)): @@ -233,7 +237,7 @@ def convert_color_space_image_tensor( elif old_color_space == ColorSpace.GRAY_ALPHA and new_color_space == ColorSpace.RGB: return _gray_to_rgb(_strip_alpha(image)) elif old_color_space == ColorSpace.GRAY_ALPHA and new_color_space == ColorSpace.RGB_ALPHA: - image, alpha = torch.tensor_split(image, indices=(-1,), dim=-3) + image, alpha = _split_alpha(image) return _add_alpha(_gray_to_rgb(image), alpha) elif old_color_space == ColorSpace.RGB and new_color_space == ColorSpace.GRAY: return _rgb_to_gray(image) @@ -244,7 +248,7 @@ def convert_color_space_image_tensor( elif old_color_space == ColorSpace.RGB_ALPHA and new_color_space == ColorSpace.GRAY: return _rgb_to_gray(_strip_alpha(image)) elif old_color_space == ColorSpace.RGB_ALPHA and new_color_space == ColorSpace.GRAY_ALPHA: - image, alpha = torch.tensor_split(image, indices=(-1,), dim=-3) + image, alpha = _split_alpha(image) return _add_alpha(_rgb_to_gray(image), alpha) elif old_color_space == ColorSpace.RGB_ALPHA and new_color_space == ColorSpace.RGB: return _strip_alpha(image)