How to draw bounding boxes on an image in PyTorch?

In this article, we are going to see how to draw bounding boxes on an image in PyTorch.

draw_bounding_boxes() method

The draw_bounding_boxes function helps us to draw bounding boxes on an image. With tensor we provide shapes in [C, H, W], where C represents the number of channels and  H, W represents the height and width respectively, this function returns an Image Tensor with bounding boxes. The bounding box should contain four points in format as (xmin, ymin, xmax, ymax), the top-left point=(xmin, ymin), and bottom-right point = (xmax, ymax). 

Syntax – torch.utils.draw_bounding_boxes(image, box)

Parameter:

• image: Tensor of shape (C x H x W) 
• box: Bounding boxes size in (xmin, ymin, xmax, ymax).

Return: Image Tensor of dtype uint8 with bounding boxes plotted.

The below image is used for demonstration:

Example 1:

The following example is to know how to change and fill colors in bounding boxes.

Python3

# Import the required libraries
import torch
import torchvision
from torchvision.io import read_image
from torchvision.utils import draw_bounding_boxes

# read input image from your computer
img = read_image('a3.png')

# bounding box are xmin, ymin, xmax, ymax
box = [330, 190, 660, 355]
box = torch.tensor(box)
box = box.unsqueeze(0)

# draw bounding box and fill color
img = draw_bounding_boxes(img, box, width=5,
                          colors="green",
                          fill=True)

# transform this image to PIL image
img = torchvision.transforms.ToPILImage()(img)

# display output
img.show()

Output:

Example 2:

The following example is to understand how to draw multiple bounding boxes on an image.

Python3

# Import the required libraries
import torch
import torchvision
from torchvision.io import read_image
from torchvision.utils import draw_bounding_boxes

# read input image from your computer
img = read_image('a3.png')

# create boxes
box_1 = [330, 220, 450, 350]
box_2 = [530, 200, 650, 320]
box = [box_1, box_2]
box = torch.tensor(box, dtype=torch.int)

# draw bounding box and fill color
img = draw_bounding_boxes(img, box, width=5, colors=[
                          "orange", "blue"], fill=True)

# transform this image to PIL image
img = torchvision.transforms.ToPILImage()(img)

# display output
img.show()

Output: