Repository: pytorch/vision
Pull request: support of float dtypes for draw_segmentation_masks
Fixes: #8138

Hey there! 

I've implemented support for draw_* methods to seamlessly handle both uint8 and float32 image types. The processed image will now be returned with the same data type as the input. Your feedback is always appreciated.

Thank you!



Problem statement:
draw_bounding_boxes and save_image  require different input dtype
Hi, when using draw_bounding boxes an image torch.tensor needs to be converted to uint8 dtype because the function expects that dtype. When inserting the image with drawn boxes to save_image, it needs to be converted back to float32 dtype. Can this be changed to easily insert the images to both functions without type conversion?

Hints from issue discussion:
Thanks for the report @SeucheAchat9115 , we should add `float32` support to the `draw_*` utils. Would you like to submit a PR for that?
I think I have no capacity to do that.

Commit message: gitignore updated for ignoring .DS_Store files

Patch to review:
diff --git a/.gitignore b/.gitignore
index f16b54061e0..91df13f6b35 100644
--- a/.gitignore
+++ b/.gitignore
@@ -42,3 +42,5 @@ xcuserdata/
 # direnv
 .direnv
 .envrc
+
+.DS_Store
\ No newline at end of file
diff --git a/test/test_utils.py b/test/test_utils.py
index 2999e84bf8a..3da87847104 100644
--- a/test/test_utils.py
+++ b/test/test_utils.py
@@ -7,8 +7,9 @@
 import numpy as np
 import pytest
 import torch
-import torchvision.transforms.functional as F
 import torchvision.utils as utils
+import torchvision.transforms.functional as F
+from torchvision.transforms.v2.functional import to_dtype
 from common_utils import assert_equal, cpu_and_cuda
 from PIL import __version__ as PILLOW_VERSION, Image, ImageColor
 
@@ -151,7 +152,7 @@ def test_draw_boxes_grayscale():
 
 def test_draw_invalid_boxes():
     img_tp = ((1, 1, 1), (1, 2, 3))
-    img_wrong1 = torch.full((3, 5, 5), 255, dtype=torch.float)
+    img_wrong1 = torch.full((3, 5, 5), 255, dtype=torch.long)
     img_wrong2 = torch.full((1, 3, 5, 5), 255, dtype=torch.uint8)
     img_correct = torch.zeros((3, 10, 10), dtype=torch.uint8)
     boxes = torch.tensor([[0, 0, 20, 20], [0, 0, 0, 0], [10, 15, 30, 35], [23, 35, 93, 95]], dtype=torch.float)
@@ -161,7 +162,7 @@ def test_draw_invalid_boxes():
 
     with pytest.raises(TypeError, match="Tensor expected"):
         utils.draw_bounding_boxes(img_tp, boxes)
-    with pytest.raises(ValueError, match="Tensor uint8 expected"):
+    with pytest.raises(ValueError, match="Tensor uint8 or float expected"):
         utils.draw_bounding_boxes(img_wrong1, boxes)
     with pytest.raises(ValueError, match="Pass individual images, not batches"):
         utils.draw_bounding_boxes(img_wrong2, boxes)
@@ -246,6 +247,25 @@ def test_draw_segmentation_masks(colors, alpha, device):
         torch.testing.assert_close(out[:, mask], interpolated_color, rtol=0.0, atol=1.0)
 
 
+def test_draw_segmentation_masks_dtypes():
+    num_masks, h, w = 2, 100, 100
+
+    masks = torch.randint(0, 2, (num_masks, h, w), dtype=torch.bool)
+
+    img_uint8 = torch.randint(0, 256, size=(3, h, w), dtype=torch.uint8)
+    out_uint8 = utils.draw_segmentation_masks(img_uint8, masks)
+
+    assert img_uint8 is not out_uint8
+    assert out_uint8.dtype == torch.uint8
+
+    img_float = to_dtype(img_uint8, torch.float, scale=True)
+    out_float = utils.draw_segmentation_masks(img_float, masks)
+
+    assert img_float is not out_float
+    assert out_float.is_floating_point ()
+
+    torch.testing.assert_close(out_uint8, to_dtype(out_float, torch.uint8, scale=True), rtol=0, atol=1.0)
+
 @pytest.mark.parametrize("device", cpu_and_cuda())
 def test_draw_segmentation_masks_errors(device):
     h, w = 10, 10
diff --git a/torchvision/utils.py b/torchvision/utils.py
index 6ec19a0e0a1..04513f42f10 100644
--- a/torchvision/utils.py
+++ b/torchvision/utils.py
@@ -164,11 +164,11 @@ def draw_bounding_boxes(
 
     """
     Draws bounding boxes on given image.
-    The values of the input image should be uint8 between 0 and 255.
+    The image values should be uint8 in [0, 255] or float in [0, 1].
     If fill is True, Resulting Tensor should be saved as PNG image.
 
     Args:
-        image (Tensor): Tensor of shape (C x H x W) and dtype uint8.
+        image (Tensor): Tensor of shape (C x H x W) and dtype uint8 or float32.
         boxes (Tensor): Tensor of size (N, 4) containing bounding boxes in (xmin, ymin, xmax, ymax) format. Note that
             the boxes are absolute coordinates with respect to the image. In other words: `0 <= xmin < xmax < W` and
             `0 <= ymin < ymax < H`.
@@ -185,15 +185,15 @@ def draw_bounding_boxes(
         font_size (int): The requested font size in points.
 
     Returns:
-        img (Tensor[C, H, W]): Image Tensor of dtype uint8 with bounding boxes plotted.
+        img (Tensor[C, H, W]): Image Tensor of dtype uint8 or float32 with bounding boxes plotted.
     """
 
     if not torch.jit.is_scripting() and not torch.jit.is_tracing():
         _log_api_usage_once(draw_bounding_boxes)
     if not isinstance(image, torch.Tensor):
         raise TypeError(f"Tensor expected, got {type(image)}")
-    elif image.dtype != torch.uint8:
-        raise ValueError(f"Tensor uint8 expected, got {image.dtype}")
+    elif not (image.dtype == torch.uint8 or image.is_floating_point()):
+        raise ValueError(f"Tensor uint8 or float expected, got {image.dtype}")
     elif image.dim() != 3:
         raise ValueError("Pass individual images, not batches")
     elif image.size(0) not in {1, 3}:
@@ -249,7 +249,7 @@ def draw_bounding_boxes(
             margin = width + 1
             draw.text((bbox[0] + margin, bbox[1] + margin), label, fill=color, font=txt_font)
 
-    return torch.from_numpy(np.array(img_to_draw)).permute(2, 0, 1).to(dtype=torch.uint8)
+    return torch.from_numpy(np.array(img_to_draw)).permute(2, 0, 1).to(dtype=image.dtype)
 
 
 @torch.no_grad()
@@ -262,10 +262,10 @@ def draw_segmentation_masks(
 
     """
     Draws segmentation masks on given RGB image.
-    The values of the input image should be uint8 between 0 and 255.
+    The image values should be uint8 in [0, 255] or float in [0, 1].
 
     Args:
-        image (Tensor): Tensor of shape (3, H, W) and dtype uint8.
+        image (Tensor): Tensor of shape (3, H, W) and dtype uint8 or float.
         masks (Tensor): Tensor of shape (num_masks, H, W) or (H, W) and dtype bool.
         alpha (float): Float number between 0 and 1 denoting the transparency of the masks.
             0 means full transparency, 1 means no transparency.
@@ -282,8 +282,8 @@ def draw_segmentation_masks(
         _log_api_usage_once(draw_segmentation_masks)
     if not isinstance(image, torch.Tensor):
         raise TypeError(f"The image must be a tensor, got {type(image)}")
-    elif image.dtype != torch.uint8:
-        raise ValueError(f"The image dtype must be uint8, got {image.dtype}")
+    elif not (image.dtype == torch.uint8 or image.is_floating_point()):
+        raise ValueError(f"The image dtype must be uint8 or float, got {image.dtype}")
     elif image.dim() != 3:
         raise ValueError("Pass individual images, not batches")
     elif image.size()[0] != 3:
@@ -297,15 +297,18 @@ def draw_segmentation_masks(
     if masks.shape[-2:] != image.shape[-2:]:
         raise ValueError("The image and the masks must have the same height and width")
 
+    original_dtype = image.dtype
+    if image.is_floating_point():
+        image = (image * 255).to(torch.uint8)
+
     num_masks = masks.size()[0]
 
     if num_masks == 0:
         warnings.warn("masks doesn't contain any mask. No mask was drawn")
         return image
 
-    out_dtype = torch.uint8
     colors = [
-        torch.tensor(color, dtype=out_dtype, device=image.device)
+        torch.tensor(color, dtype=torch.uint8, device=image.device)
         for color in _parse_colors(colors, num_objects=num_masks)
     ]
 
@@ -315,7 +318,10 @@ def draw_segmentation_masks(
         img_to_draw[:, mask] = color[:, None]
 
     out = image * (1 - alpha) + img_to_draw * alpha
-    return out.to(out_dtype)
+    if original_dtype in {torch.float16, torch.float32, torch.float64}:
+        out = out.float() / 255.0
+
+    return out.to(original_dtype)
 
 
 @torch.no_grad()
@@ -330,10 +336,10 @@ def draw_keypoints(
 
     """
     Draws Keypoints on given RGB image.
-    The values of the input image should be uint8 between 0 and 255.
+    The image values should be uint8 in [0, 255] or float in [0, 1].
 
     Args:
-        image (Tensor): Tensor of shape (3, H, W) and dtype uint8.
+        image (Tensor): Tensor of shape (3, H, W) and dtype uint8 or float.
         keypoints (Tensor): Tensor of shape (num_instances, K, 2) the K keypoints location for each of the N instances,
             in the format [x, y].
         connectivity (List[Tuple[int, int]]]): A List of tuple where,
@@ -344,15 +350,15 @@ def draw_keypoints(
         width (int): Integer denoting width of line connecting keypoints.
 
     Returns:
-        img (Tensor[C, H, W]): Image Tensor of dtype uint8 with keypoints drawn.
+        img (Tensor[C, H, W]): Image Tensor of dtype uint8 or float32 with keypoints drawn.
     """
 
     if not torch.jit.is_scripting() and not torch.jit.is_tracing():
         _log_api_usage_once(draw_keypoints)
     if not isinstance(image, torch.Tensor):
         raise TypeError(f"The image must be a tensor, got {type(image)}")
-    elif image.dtype != torch.uint8:
-        raise ValueError(f"The image dtype must be uint8, got {image.dtype}")
+    elif not (image.dtype == torch.uint8 or image.is_floating_point()):
+        raise ValueError(f"The image dtype must be uint8 or float, got {image.dtype}")
     elif image.dim() != 3:
         raise ValueError("Pass individual images, not batches")
     elif image.size()[0] != 3:
@@ -387,7 +393,7 @@ def draw_keypoints(
                     width=width,
                 )
 
-    return torch.from_numpy(np.array(img_to_draw)).permute(2, 0, 1).to(dtype=torch.uint8)
+    return torch.from_numpy(np.array(img_to_draw)).permute(2, 0, 1).to(dtype=image.dtype)
 
 
 # Flow visualization code adapted from https://github.com/tomrunia/OpticalFlow_Visualization


Code review comments:
diff_hunk: @@ -297,15 +297,18 @@ def draw_segmentation_masks(
     if masks.shape[-2:] != image.shape[-2:]:
         raise ValueError("The image and the masks must have the same height and width")
 
+    original_dtype = image.dtype
+    if image.is_floating_point():
+        image = (image * 255).to(torch.uint8)
path: torchvision/utils.py
text: @user1:
Let's use `to_dtype` here instead, like you did in the test.

diff_hunk: @@ -315,7 +318,10 @@ def draw_segmentation_masks(
         img_to_draw[:, mask] = color[:, None]
 
     out = image * (1 - alpha) + img_to_draw * alpha
-    return out.to(out_dtype)
+    if original_dtype in {torch.float16, torch.float32, torch.float64}:
path: torchvision/utils.py
text: @user1:
Check for `is_floating_point()` instead.

diff_hunk: @@ -315,7 +318,10 @@ def draw_segmentation_masks(
         img_to_draw[:, mask] = color[:, None]
 
     out = image * (1 - alpha) + img_to_draw * alpha
-    return out.to(out_dtype)
+    if original_dtype in {torch.float16, torch.float32, torch.float64}:
+        out = out.float() / 255.0
path: torchvision/utils.py
text: @user1:
Same: Let's use `to_dtype` here instead, like you did in the test.

diff_hunk: @@ -246,6 +247,25 @@ def test_draw_segmentation_masks(colors, alpha, device):
         torch.testing.assert_close(out[:, mask], interpolated_color, rtol=0.0, atol=1.0)
 
 
+def test_draw_segmentation_masks_dtypes():
+    num_masks, h, w = 2, 100, 100
+
+    masks = torch.randint(0, 2, (num_masks, h, w), dtype=torch.bool)
+
+    img_uint8 = torch.randint(0, 256, size=(3, h, w), dtype=torch.uint8)
+    out_uint8 = utils.draw_segmentation_masks(img_uint8, masks)
+
+    assert img_uint8 is not out_uint8
+    assert out_uint8.dtype == torch.uint8
+
+    img_float = to_dtype(img_uint8, torch.float, scale=True)
+    out_float = utils.draw_segmentation_masks(img_float, masks)
+
+    assert img_float is not out_float
+    assert out_float.is_floating_point ()
+
+    torch.testing.assert_close(out_uint8, to_dtype(out_float, torch.uint8, scale=True), rtol=0, atol=1.0)
path: test/test_utils.py
text: @user1:
```suggestion
    torch.testing.assert_close(out_uint8, to_dtype(out_float, torch.uint8, scale=True), rtol=0, atol=1)
```

diff_hunk: @@ -246,6 +247,25 @@ def test_draw_segmentation_masks(colors, alpha, device):
         torch.testing.assert_close(out[:, mask], interpolated_color, rtol=0.0, atol=1.0)
 
 
+def test_draw_segmentation_masks_dtypes():
+    num_masks, h, w = 2, 100, 100
+
+    masks = torch.randint(0, 2, (num_masks, h, w), dtype=torch.bool)
+
+    img_uint8 = torch.randint(0, 256, size=(3, h, w), dtype=torch.uint8)
+    out_uint8 = utils.draw_segmentation_masks(img_uint8, masks)
+
+    assert img_uint8 is not out_uint8
+    assert out_uint8.dtype == torch.uint8
+
+    img_float = to_dtype(img_uint8, torch.float, scale=True)
+    out_float = utils.draw_segmentation_masks(img_float, masks)
+
+    assert img_float is not out_float
+    assert out_float.is_floating_point ()
path: test/test_utils.py
text: @user1:
```suggestion
    assert out_float.is_floating_point()
```

diff_hunk: @@ -42,3 +42,5 @@ xcuserdata/
 # direnv
 .direnv
 .envrc
+
+.DS_Store
path: .gitignore
text: @user1:
Let's revert these changes, we shouldn't need to modify this file here

```suggestion

```