Add utility to draw bounding boxes

docs/source/utils.rst

@@ -7,3 +7,4 @@ torchvision.utils
 
 .. autofunction:: save_image
 
+.. autofunction:: draw_bounding_boxes

test/test_utils.py

@@ -6,7 +6,7 @@
 import unittest
 from io import BytesIO
 import torchvision.transforms.functional as F
-from PIL import Image
+from PIL import Image, ImageDraw
 
 
 class Tester(unittest.TestCase):
@@ -79,6 +79,15 @@ def test_save_image_single_pixel_file_object(self):
             self.assertTrue(torch.equal(F.to_tensor(img_orig), F.to_tensor(img_bytes)),
                             'Pixel Image not stored in file object')
 
+    def test_draw_boxes(self):
+        img = torch.randint(0, 255, (3, 226, 226), dtype=torch.uint8)
+        boxes = torch.tensor([[0, 0, 100, 100], [0, 0, 0, 0],
+                             [10, 15, 30, 35], [23, 35, 93, 95]], dtype=torch.float)
+        labels = ['a', 'b', 'c', 'd']
+        utils.draw_bounding_boxes(img, boxes, labels=labels)
+
+        return True
+
 
 if __name__ == '__main__':
     unittest.main()

torchvision/utils.py

@@ -1,8 +1,13 @@
 from typing import Union, Optional, List, Tuple, Text, BinaryIO
-import io
 import pathlib
 import torch
 import math
+import numpy as np
+from PIL import Image, ImageDraw
+from PIL.ImageFont import ImageFont
+
+__all__ = ["make_grid", "save_image", "draw_bounding_boxes"]
+
 irange = range
 
 
@@ -121,10 +126,60 @@ def save_image(
             If a file object was used instead of a filename, this parameter should always be used.
         **kwargs: Other arguments are documented in ``make_grid``.
     """
-    from PIL import Image
     grid = make_grid(tensor, nrow=nrow, padding=padding, pad_value=pad_value,
                      normalize=normalize, range=range, scale_each=scale_each)
     # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
     ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
     im = Image.fromarray(ndarr)
     im.save(fp, format=format)
+
+
+def draw_bounding_boxes(
+    image: torch.Tensor,
+    boxes: torch.Tensor,
+    colors: Optional[List[str]] = None,
+    labels: Optional[List[str]] = None,
+    width: int = 1,
+    font: Optional[ImageFont] = None
+) -> torch.Tensor:
+
+    """
+    Draws bounding boxes on given image.
+    The values of the input image should be uint8 between 0 and 255.
+
+    Args:
+        image (Tensor): Tensor of shape (C x H x W)
+        bboxes (Tensor): Tensor of size (N, 4) containing bounding boxes in (xmin, ymin, xmax, ymax) format.
+        colors (List): List containing the colors of bounding boxes excluding background.
+        labels (List): List containing the labels of bounding boxes excluding background.
+        width (int): Width of bounding box.
+        font (ImageFont): The PIL ImageFont object used to for drawing the labels.
+    """
+
+    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 image.dim() != 3:
+        raise ValueError("Pass individual images, not batches")
+
+    if image.requires_grad:
+        image = image.detach()
+    if boxes.requires_grad:
+        boxes = boxes.detach()
+
+    ndarr = image.permute(1, 2, 0).numpy()
+    img_to_draw = Image.fromarray(ndarr)
+
+    img_boxes = boxes.to(torch.int64).tolist()
+
+    draw = ImageDraw.Draw(img_to_draw)
+
+    for i, bbox in enumerate(img_boxes):
+        color = None if colors is None else colors[i]
+        draw.rectangle(bbox, width=width, outline=color)
+
+        if labels is not None:
+            draw.text((bbox[0], bbox[1]), labels[i], fill=color, font=font)
+
+    return torch.from_numpy(np.array(img_to_draw)).permute(2, 0, 1)