How to draw bounding boxes on an image in PyTorch?
Last Updated :
22 Apr, 2022
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 torch
import torchvision
from torchvision.io import read_image
from torchvision.utils import draw_bounding_boxes
img = read_image( 'a3.png' )
box = [ 330 , 190 , 660 , 355 ]
box = torch.tensor(box)
box = box.unsqueeze( 0 )
img = draw_bounding_boxes(img, box, width = 5 ,
colors = "green" ,
fill = True )
img = torchvision.transforms.ToPILImage()(img)
img.show()
|
Output:
Example 2:
The following example is to understand how to draw multiple bounding boxes on an image.
Python3
import torch
import torchvision
from torchvision.io import read_image
from torchvision.utils import draw_bounding_boxes
img = read_image( 'a3.png' )
box_1 = [ 330 , 220 , 450 , 350 ]
box_2 = [ 530 , 200 , 650 , 320 ]
box = [box_1, box_2]
box = torch.tensor(box, dtype = torch. int )
img = draw_bounding_boxes(img, box, width = 5 , colors = [
"orange" , "blue" ], fill = True )
img = torchvision.transforms.ToPILImage()(img)
img.show()
|
Output:
Like Article
Suggest improvement
Share your thoughts in the comments
Please Login to comment...