Repository: pytorch/vision

Issue:
[RFC] How do we want to deal with images that include alpha channels?
This discussion started in https://github.com/pytorch/vision/pull/5500#discussion_r816503203 and @vfdev-5 and I continued offline.

PIL as well as our image reading functions support RGBA images

https://github.com/pytorch/vision/blob/95d418970e6dbf2e4d928a204c4e620da7bccdc0/torchvision/io/image.py#L16-L31

but our color transformations currently only support RGB images ignoring an extra alpha channel. This leads to wrong results. One thing that we agreed upon is that these transforms should fail if anything but 3 channels is detected.


Still, some datasets include non-RGB images so we need to deal with this for a smooth UX. Previously we implicitly converted every image to RGB before returning it from a dataset

https://github.com/pytorch/vision/blob/f9fbc104c02f277f9485d9f8727f3d99a1cf5f0b/torchvision/datasets/folder.py#L245-L249

Since we no longer decode images in the datasets, we need to provide a solution for the users here. I currently see two possible options:

1. We could deal with this on a per-image basis within the dataset. For example, the train split of ImageNet contains a single RGBA image. We could simply perform an appropriate conversion for irregular image modes in the dataset so this issue is abstracted away from the user. `tensorflow-datasets` uses this approach: https://github.com/tensorflow/datasets/blob/a1caff379ed3164849fdefd147473f72a22d3fa7/tensorflow_datasets/image_classification/imagenet.py#L105-L131
2. The most common non-RGB image in datasets are grayscale images. For example, the train split of ImageNet contains 19970 grayscale images. Thus, the users will need a `transforms.ConvertImageColorSpace("rgb")` in most cases anyway. If that would support RGBA to RGB conversions the problem would also be solved. The conversion happens with this formula:

    ```
    pixel_new = (1 - alpha) * background + alpha * pixel_old
    ```
    
    where `pixel_{old|new}` is a single value from a color channel. Since we don't know `background` we need to either make assumptions or require the user to provide a value for it. I'd wager a guess that in 99% of the cases the background is white. i.e. `background == 1`, but we can't be sure about that.
    
    Another issue with this is that the user has no option to set the background on a per-image basis in the transforms pipeline if that is needed.

    In special case for `alpha == 1` everywhere, the equation above simplifies to

    ```
    pixel_new = pixel_old
    ```

    which is equivalent to stripping the alpha channel. We could check for that and only perform the RGBA to RGB transform if the condition holds or the user supplies a background color.




cc @pmeier @vfdev-5 @datumbox @bjuncek


Discussion:
I personally would go with a combined approach:

1. If we are aware of any images in a dataset outside of what we can convert automatically without assumptions, the dataset should handle these images.
2. Add a `background` parameter to `convert_image_color_space`. If this is not passed and the alpha channel is not 1 everywhere fail. 
I agree that for torchvision supported datasets if alpha channel is not used later by a CNN then we can strip it from the beginning. For example, as you say as ImageNet can have RGBA images, it would be better to provide only RGB images as it is done right now in the stable API. One way I could think about that is to have a preset datapipe transforms that could be deactivated:
```python
dp = datasets.load("imagenet", decode_presets=datasets.imagenet.decode_preset)
```

When you say that "image reading functions support RGBA images". RGBA is a conventional as PIL "RGBA" mode or premultiplied as PIL "RGBa" ?  (for differences between: https://shawnhargreaves.com/blog/premultiplied-alpha.html)

Another question, `features.Image` can have ColorSpace as RGBA ?

Some thoughts on `ConvertImageColorSpace("rgb")`. This means that we should be aware of the current mode and 
possibly can transform to hsv, hsl and other modes (e.g. like opencv).

By the way, PIL handles correctly alpha channel on transformations, so impacted color transforms are only related to `features.Image`.

As for choosing `background` as 1, 255 or white, a problem can be seen on geometric transforms where default fill color is 0 or black, thus either user have to add everywhere fill color as white or they will have 2 types of no-data pixels: black and white...


> possibly can transform to hsv, hsl and other modes

I don't think that's the focus right now. Maybe on the long long future but definitely not something we plan to do anytime soon.

> As for choosing background as 1, 255 or white, .....

+1 on that. Given that our existing functional transforms have fill=0, does it make sense to go with black?


@vfdev-5 

> I agree that for torchvision supported datasets if alpha channel is not used later by a CNN then we can strip it from the beginning. 

I might be pedantic here, but I want to be sure we are all on the same page: we cannot simply strip the alpha channel unless the alpha channel is all ones or 255s depending on the input type. If that is not the case, stripping the alpha channel has an effect on the actual color of the image as can be seen in the equation above.

> For example, as you say as ImageNet can have RGBA images, it would be better to provide only RGB images as it is done right now in the stable API. One way I could think about that is to have a preset datapipe transforms that could be deactivated:

How would this approach know which image needs the special conversion? If we have a general approach like this, we also need to solve the conversion in a general sense.

> When you say that "image reading functions support RGBA images". RGBA is a conventional as PIL "RGBA" mode or premultiplied as PIL "RGBa" ?

I don't know. @datumbox has introduced the enum. Do you have any insights here?

> Another question, `features.Image` can have ColorSpace as RGBA ?

Yes. For now, we only support 

https://github.com/pytorch/vision/blob/71d2bb0bc67044f55d38bfddf04e05be0343deab/torchvision/prototype/features/_image.py#L16-L19

But if we decide how we want to handle RGBA we can of course add it.

> Some thoughts on `ConvertImageColorSpace("rgb")`. This means that we should be aware of the current mode and

That is already implemented. `features.Image`'s and `PIL.Image.Image`'s store this information. Only for vanilla tensors, the user needs to provide it manually:

https://github.com/pytorch/vision/blob/71d2bb0bc67044f55d38bfddf04e05be0343deab/torchvision/prototype/transforms/_meta.py#L55-L78

> possibly can transform to hsv, hsl and other modes (e.g. like opencv).

I agree with @datumbox here. While the API allows all these conversions, we shouldn't add anything here unless there is a specific user request. If we have other color spaces that are equivalent to RGB such as HSV or HSL, we definitely need a `color_space` parameter on all color transforms and convert to RGB before we start applying the transform.

> As for choosing `background` as 1, 255 or white, a problem can be seen on geometric transforms where default fill color is 0 or black, thus either user have to add everywhere fill color as white or they will have 2 types of no-data pixels: black and white...

The alpha channel is different from a fill color. Choosing alpha all white means no transparency and thus the background is irrelevant. For our all of our current datasets, RGBA images are outliers that are most likely not intended by the authors. They happen if the images are scraped from the internet and the person who uploaded them made a mistake to include an alpha channel  although there was no need for it. For example, this is the RGBA image from ImageNet (in thumbnail size to limit copyright implications) *including* the alpha channel:

![imagenet_rgba](https://user-images.githubusercontent.com/6849766/156401155-3c1d61ad-fb48-47f8-bea1-f3ee9f636ac4.png)




> How would this approach know which image needs the special conversion? If we have a general approach like this, we also need to solve the conversion in a general sense.

Maybe, I'm missing something important, but for me, everything could just follow the same logic as actual API does with Pillow images: when image is read and decoded you'll have a color space. Prior info is that for imagenet we provide RGB only images. So, all images with non `RGB` space should be converted to `RGB`.


> So, all images with non `RGB` space should be converted to `RGB`.

That is the crucial point I was touching on in this issue. How are we going to implement the conversion from RGBA to RGB? Not all image files carry information about the background color in the first place. The image above is one of them. In such a case [PIL falls back](https://github.com/python-pillow/Pillow/blob/92c26a77ca53a2bfbd8804f009c6c8755d0e5a43/src/PIL/Image.py#L1056-L1061) to [Floyd-Steinberg dithering](https://en.wikipedia.org/wiki/Floyd%E2%80%93Steinberg_dithering).

Not only would we need to implement custom kernels for that, but dithering might also significantly reduce the image quality. Compare this [test image](https://www.w3.org/Graphics/PNG/Alphatest.html)

![](https://www.w3.org/Graphics/PNG/text2.png)

`.convert("RGB")`'ed by PIL with Floyd-Steinberg

![dithered](https://user-images.githubusercontent.com/6849766/156438720-797b0e12-1f77-4317-af6b-cc33dbe8f123.png)

or converted by the formula above assuming a white background

![converted](https://user-images.githubusercontent.com/6849766/156439676-cb90e9e2-9799-4329-8f7f-082f97043804.png)


So the two key approaches discussed here are:
1. Handle RGBA to RGB conversion exactly like PIL using Floyd-Steinberg
2. Use the formula provided by Philip and assume white background

What is the approach used by CV2?
I'll dig in what PIL is exactly doing, because what I failed to realize is that they also seem to assume a white background for the dithering.
PIL [assumes a white background](https://github.com/python-pillow/Pillow/blob/c58d2817bc891c26e6b8098b8909c0eb2e7ce61b/src/libImaging/Convert.c#L421-L431).
@pmeier Can you check also CV2? If they do the same, I think it would make sense to make the same assumption on TorchVision. 

@vfdev-5 Thoughts?
### CV2

I'm not 100% sure, but my best guess is that white is also assumed. I've tracked it to this [OpenCL call](https://github.com/opencv/opencv/blob/17234f82d025e3bbfbf611089637e5aa2038e7b8/modules/imgproc/src/color_rgb.dispatch.cpp#L426-L437). There are multiple issues on the OpenCV repo such as opencv/opencv#13135 that report that the conversion does not take alpha into account. I think the consensus is that one should not use `cv2.cvtColor` for RGBA -> RGB conversions, but rather blend if the alpha actually contains information. Ignoring the alpha channel is effectively the same than setting it to white.

### skimage

skimage has an `background` parameter on their [`rgba2rgb()` conversion function](https://scikit-image.org/docs/dev/api/skimage.color.html#rgba2rgb). It defaults to white.

Thus, I think we are fairly safe if we also assume white and provide an option to set it manually if required.
After some offline discussion, we decided to align with PIL for now. The only difference should be that we should fail the transformation if the alpha channel is not the max value everywhere. This way we can implement the correct conversion as detailed in my top comment later without worrying about BC.
> PIL [assumes a white background](https://github.com/python-pillow/Pillow/blob/c58d2817bc891c26e6b8098b8909c0eb2e7ce61b/src/libImaging/Convert.c#L421-L431).

@pmeier You mean rgba2rgb method: https://github.com/python-pillow/Pillow/blob/c58d2817bc891c26e6b8098b8909c0eb2e7ce61b/src/libImaging/Convert.c#L443-L453 right ?

How do you see that it assumes white background ?
From the implementation we can say that pointers `in` and `out` has 4 values per pixel. By the last `in++` we ignore the forth value and by `*out++ = 255;` we set output forth value as 255.

Pixel size is 4 for RGB: https://github.com/python-pillow/Pillow/blob/95cff6e959bb3c37848158ed2145d49d49806a31/src/libImaging/Storage.c#L125-L129



> In such a case [PIL falls back](https://github.com/python-pillow/Pillow/blob/92c26a77ca53a2bfbd8804f009c6c8755d0e5a43/src/PIL/Image.py#L1056-L1061) to [Floyd-Steinberg dithering](https://en.wikipedia.org/wiki/Floyd%E2%80%93Steinberg_dithering).

In PIL docs they say to use dithering for conversion like RGB to P or 1 
https://github.com/python-pillow/Pillow/blob/92c26a77ca53a2bfbd8804f009c6c8755d0e5a43/src/PIL/Image.py#L921-L922
same for the code:

https://github.com/python-pillow/Pillow/blob/c58d2817bc891c26e6b8098b8909c0eb2e7ce61b/src/libImaging/Convert.c#L1572-L1574

https://github.com/python-pillow/Pillow/blob/c58d2817bc891c26e6b8098b8909c0eb2e7ce61b/src/libImaging/Convert.c#L1608-L1610
You are right. PIL also seems to only "strip" the alpha channel by setting it to 255 everywhere:

```py
import PIL.Image
import torch
from torchvision.transforms.functional import pil_to_tensor

image = PIL.Image.open("text2.png")
a = pil_to_tensor(image)[:3, ...]
b = pil_to_tensor(image.convert("RGB"))

torch.testing.assert_close(a, b)
```

Thus, the behavior is aligned with OpenCV. This doesn't change the plan from https://github.com/pytorch/vision/issues/5510#issuecomment-1061529561.

Resolving patch:
diff --git a/torchvision/prototype/features/_image.py b/torchvision/prototype/features/_image.py
--- a/torchvision/prototype/features/_image.py
+++ b/torchvision/prototype/features/_image.py
@@ -15,8 +15,23 @@
 
 class ColorSpace(StrEnum):
     OTHER = StrEnum.auto()
-    GRAYSCALE = StrEnum.auto()
+    GRAY = StrEnum.auto()
+    GRAY_ALPHA = StrEnum.auto()
     RGB = StrEnum.auto()
+    RGB_ALPHA = StrEnum.auto()
+
+    @classmethod
+    def from_pil_mode(cls, mode: str) -> ColorSpace:
+        if mode == "L":
+            return cls.GRAY
+        elif mode == "LA":
+            return cls.GRAY_ALPHA
+        elif mode == "RGB":
+            return cls.RGB
+        elif mode == "RGBA":
+            return cls.RGB_ALPHA
+        else:
+            return cls.OTHER
 
 
 class Image(_Feature):
@@ -71,13 +86,17 @@ def guess_color_space(data: torch.Tensor) -> ColorSpace:
         if data.ndim < 2:
             return ColorSpace.OTHER
         elif data.ndim == 2:
-            return ColorSpace.GRAYSCALE
+            return ColorSpace.GRAY
 
         num_channels = data.shape[-3]
         if num_channels == 1:
-            return ColorSpace.GRAYSCALE
+            return ColorSpace.GRAY
+        elif num_channels == 2:
+            return ColorSpace.GRAY_ALPHA
         elif num_channels == 3:
             return ColorSpace.RGB
+        elif num_channels == 4:
+            return ColorSpace.RGB_ALPHA
         else:
             return ColorSpace.OTHER
 
diff --git a/torchvision/prototype/transforms/_meta.py b/torchvision/prototype/transforms/_meta.py
--- a/torchvision/prototype/transforms/_meta.py
+++ b/torchvision/prototype/transforms/_meta.py
@@ -48,11 +48,11 @@ def __init__(
         super().__init__()
 
         if isinstance(color_space, str):
-            color_space = features.ColorSpace[color_space]
+            color_space = features.ColorSpace.from_str(color_space)
         self.color_space = color_space
 
         if isinstance(old_color_space, str):
-            old_color_space = features.ColorSpace[old_color_space]
+            old_color_space = features.ColorSpace.from_str(old_color_space)
         self.old_color_space = old_color_space
 
     def _transform(self, input: Any, params: Dict[str, Any]) -> Any:
@@ -72,13 +72,6 @@ def _transform(self, input: Any, params: Dict[str, Any]) -> Any:
                 input, old_color_space=self.old_color_space, new_color_space=self.color_space
             )
         elif isinstance(input, PIL.Image.Image):
-            old_color_space = {
-                "L": features.ColorSpace.GRAYSCALE,
-                "RGB": features.ColorSpace.RGB,
-            }.get(input.mode, features.ColorSpace.OTHER)
-
-            return F.convert_image_color_space_pil(
-                input, old_color_space=old_color_space, new_color_space=self.color_space
-            )
+            return F.convert_image_color_space_pil(input, color_space=self.color_space)
         else:
             return input
diff --git a/torchvision/prototype/transforms/functional/_meta.py b/torchvision/prototype/transforms/functional/_meta.py
--- a/torchvision/prototype/transforms/functional/_meta.py
+++ b/torchvision/prototype/transforms/functional/_meta.py
@@ -1,9 +1,10 @@
+from typing import Tuple, Optional
+
 import PIL.Image
 import torch
 from torchvision.prototype.features import BoundingBoxFormat, ColorSpace
 from torchvision.transforms import functional_tensor as _FT, functional_pil as _FP
 
-
 get_dimensions_image_tensor = _FT.get_dimensions
 get_dimensions_image_pil = _FP.get_dimensions
 
@@ -57,41 +58,88 @@ def convert_bounding_box_format(
     return bounding_box
 
 
-def _grayscale_to_rgb_tensor(grayscale: torch.Tensor) -> torch.Tensor:
-    repeats = [1] * grayscale.ndim
-    repeats[-3] = 3
-    return grayscale.repeat(repeats)
+def _split_alpha(image: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    return image[..., :-1, :, :], image[..., -1:, :, :]
 
 
-def convert_image_color_space_tensor(
-    image: torch.Tensor, old_color_space: ColorSpace, new_color_space: ColorSpace
-) -> torch.Tensor:
-    if new_color_space == old_color_space:
-        return image.clone()
+def _strip_alpha(image: torch.Tensor) -> torch.Tensor:
+    image, alpha = _split_alpha(image)
+    if not torch.all(alpha == _FT._max_value(alpha.dtype)):
+        raise RuntimeError(
+            "Stripping the alpha channel if it contains values other than the max value is not supported."
+        )
+    return image
 
-    if old_color_space == ColorSpace.GRAYSCALE:
-        image = _grayscale_to_rgb_tensor(image)
 
-    if new_color_space == ColorSpace.GRAYSCALE:
-        image = _FT.rgb_to_grayscale(image)
+def _add_alpha(image: torch.Tensor, alpha: Optional[torch.Tensor] = None) -> torch.Tensor:
+    if alpha is None:
+        shape = list(image.shape)
+        shape[-3] = 1
+        alpha = torch.full(shape, _FT._max_value(image.dtype), dtype=image.dtype, device=image.device)
+    return torch.cat((image, alpha), dim=-3)
 
-    return image
 
-
-def _grayscale_to_rgb_pil(grayscale: PIL.Image.Image) -> PIL.Image.Image:
-    return grayscale.convert("RGB")
+def _gray_to_rgb(grayscale: torch.Tensor) -> torch.Tensor:
+    repeats = [1] * grayscale.ndim
+    repeats[-3] = 3
+    return grayscale.repeat(repeats)
 
 
-def convert_image_color_space_pil(
-    image: PIL.Image.Image, old_color_space: ColorSpace, new_color_space: ColorSpace
-) -> PIL.Image.Image:
-    if new_color_space == old_color_space:
-        return image.copy()
+_rgb_to_gray = _FT.rgb_to_grayscale
 
-    if old_color_space == ColorSpace.GRAYSCALE:
-        image = _grayscale_to_rgb_pil(image)
 
-    if new_color_space == ColorSpace.GRAYSCALE:
-        image = _FP.to_grayscale(image)
+def convert_image_color_space_tensor(
+    image: torch.Tensor, old_color_space: ColorSpace, new_color_space: ColorSpace
+) -> torch.Tensor:
+    if new_color_space == old_color_space:
+        return image.clone()
 
-    return image
+    if old_color_space == ColorSpace.OTHER or new_color_space == ColorSpace.OTHER:
+        raise RuntimeError(f"Conversion to or from {ColorSpace.OTHER} is not supported.")
+
+    if old_color_space == ColorSpace.GRAY and new_color_space == ColorSpace.GRAY_ALPHA:
+        return _add_alpha(image)
+    elif old_color_space == ColorSpace.GRAY and new_color_space == ColorSpace.RGB:
+        return _gray_to_rgb(image)
+    elif old_color_space == ColorSpace.GRAY and new_color_space == ColorSpace.RGB_ALPHA:
+        return _add_alpha(_gray_to_rgb(image))
+    elif old_color_space == ColorSpace.GRAY_ALPHA and new_color_space == ColorSpace.GRAY:
+        return _strip_alpha(image)
+    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 = _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)
+    elif old_color_space == ColorSpace.RGB and new_color_space == ColorSpace.GRAY_ALPHA:
+        return _add_alpha(_rgb_to_gray(image))
+    elif old_color_space == ColorSpace.RGB and new_color_space == ColorSpace.RGB_ALPHA:
+        return _add_alpha(image)
+    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 = _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)
+    else:
+        raise RuntimeError(f"Conversion from {old_color_space} to {new_color_space} is not supported.")
+
+
+_COLOR_SPACE_TO_PIL_MODE = {
+    ColorSpace.GRAY: "L",
+    ColorSpace.GRAY_ALPHA: "LA",
+    ColorSpace.RGB: "RGB",
+    ColorSpace.RGB_ALPHA: "RGBA",
+}
+
+
+def convert_image_color_space_pil(image: PIL.Image.Image, color_space: ColorSpace) -> PIL.Image.Image:
+    old_mode = image.mode
+    try:
+        new_mode = _COLOR_SPACE_TO_PIL_MODE[color_space]
+    except KeyError:
+        raise ValueError(f"Conversion from {ColorSpace.from_pil_mode(old_mode)} to {color_space} is not supported.")
+
+    return image.convert(new_mode)


Test patch:
diff --git a/test/test_prototype_transforms.py b/test/test_prototype_transforms.py
--- a/test/test_prototype_transforms.py
+++ b/test/test_prototype_transforms.py
@@ -102,7 +102,14 @@ def test_mixup_cutmix(self, transform, input):
             (
                 transform,
                 itertools.chain.from_iterable(
-                    fn(dtypes=[torch.uint8], extra_dims=[(4,)])
+                    fn(
+                        color_spaces=[
+                            features.ColorSpace.GRAY,
+                            features.ColorSpace.RGB,
+                        ],
+                        dtypes=[torch.uint8],
+                        extra_dims=[(4,)],
+                    )
                     for fn in [
                         make_images,
                         make_vanilla_tensor_images,
@@ -152,3 +159,32 @@ def test_normalize(self, transform, input):
     )
     def test_random_resized_crop(self, transform, input):
         transform(input)
+
+    @parametrize(
+        [
+            (
+                transforms.ConvertImageColorSpace(color_space=new_color_space, old_color_space=old_color_space),
+                itertools.chain.from_iterable(
+                    [
+                        fn(color_spaces=[old_color_space])
+                        for fn in (
+                            make_images,
+                            make_vanilla_tensor_images,
+                            make_pil_images,
+                        )
+                    ]
+                ),
+            )
+            for old_color_space, new_color_space in itertools.product(
+                [
+                    features.ColorSpace.GRAY,
+                    features.ColorSpace.GRAY_ALPHA,
+                    features.ColorSpace.RGB,
+                    features.ColorSpace.RGB_ALPHA,
+                ],
+                repeat=2,
+            )
+        ]
+    )
+    def test_convert_image_color_space(self, transform, input):
+        transform(input)
diff --git a/test/test_prototype_transforms_functional.py b/test/test_prototype_transforms_functional.py
--- a/test/test_prototype_transforms_functional.py
+++ b/test/test_prototype_transforms_functional.py
@@ -7,33 +7,44 @@
 from torch import jit
 from torch.nn.functional import one_hot
 from torchvision.prototype import features
+from torchvision.transforms.functional_tensor import _max_value as get_max_value
 
 make_tensor = functools.partial(torch.testing.make_tensor, device="cpu")
 
 
-def make_image(size=None, *, color_space, extra_dims=(), dtype=torch.float32):
+def make_image(size=None, *, color_space, extra_dims=(), dtype=torch.float32, constant_alpha=True):
     size = size or torch.randint(16, 33, (2,)).tolist()
 
-    num_channels = {
-        features.ColorSpace.GRAYSCALE: 1,
-        features.ColorSpace.RGB: 3,
-    }[color_space]
+    try:
+        num_channels = {
+            features.ColorSpace.GRAY: 1,
+            features.ColorSpace.GRAY_ALPHA: 2,
+            features.ColorSpace.RGB: 3,
+            features.ColorSpace.RGB_ALPHA: 4,
+        }[color_space]
+    except KeyError as error:
+        raise pytest.UsageError() from error
 
     shape = (*extra_dims, num_channels, *size)
-    if dtype.is_floating_point:
-        data = torch.rand(shape, dtype=dtype)
-    else:
-        data = torch.randint(0, torch.iinfo(dtype).max, shape, dtype=dtype)
+    max_value = get_max_value(dtype)
+    data = make_tensor(shape, low=0, high=max_value, dtype=dtype)
+    if color_space in {features.ColorSpace.GRAY_ALPHA, features.ColorSpace.RGB_ALPHA} and constant_alpha:
+        data[..., -1, :, :] = max_value
     return features.Image(data, color_space=color_space)
 
 
-make_grayscale_image = functools.partial(make_image, color_space=features.ColorSpace.GRAYSCALE)
+make_grayscale_image = functools.partial(make_image, color_space=features.ColorSpace.GRAY)
 make_rgb_image = functools.partial(make_image, color_space=features.ColorSpace.RGB)
 
 
 def make_images(
     sizes=((16, 16), (7, 33), (31, 9)),
-    color_spaces=(features.ColorSpace.GRAYSCALE, features.ColorSpace.RGB),
+    color_spaces=(
+        features.ColorSpace.GRAY,
+        features.ColorSpace.GRAY_ALPHA,
+        features.ColorSpace.RGB,
+        features.ColorSpace.RGB_ALPHA,
+    ),
     dtypes=(torch.float32, torch.uint8),
     extra_dims=((4,), (2, 3)),
 ):
@@ -48,15 +59,12 @@ def randint_with_tensor_bounds(arg1, arg2=None, **kwargs):
     low, high = torch.broadcast_tensors(
         *[torch.as_tensor(arg) for arg in ((0, arg1) if arg2 is None else (arg1, arg2))]
     )
-    try:
-        return torch.stack(
-            [
-                torch.randint(low_scalar, high_scalar, (), **kwargs)
-                for low_scalar, high_scalar in zip(low.flatten().tolist(), high.flatten().tolist())
-            ]
-        ).reshape(low.shape)
-    except RuntimeError as error:
-        raise error
+    return torch.stack(
+        [
+            torch.randint(low_scalar, high_scalar, (), **kwargs)
+            for low_scalar, high_scalar in zip(low.flatten().tolist(), high.flatten().tolist())
+        ]
+    ).reshape(low.shape)
 
 
 def make_bounding_box(*, format, image_size=(32, 32), extra_dims=(), dtype=torch.int64):
@@ -83,8 +91,8 @@ def make_bounding_box(*, format, image_size=(32, 32), extra_dims=(), dtype=torch
         w = randint_with_tensor_bounds(1, torch.minimum(cx, width - cx) + 1)
         h = randint_with_tensor_bounds(1, torch.minimum(cy, width - cy) + 1)
         parts = (cx, cy, w, h)
-    else:  # format == features.BoundingBoxFormat._SENTINEL:
-        raise ValueError()
+    else:
+        raise pytest.UsageError()
 
     return features.BoundingBox(torch.stack(parts, dim=-1).to(dtype), format=format, image_size=image_size)