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:
