¿Cómo dibujar cuadros delimitadores en una imagen en PyTorch?

En este artículo, vamos a ver cómo dibujar cuadros delimitadores en una imagen en PyTorch.

método draw_bounding_boxes()

La función draw_bounding_boxes nos ayuda a dibujar cuadros delimitadores en una imagen. Con tensor proporcionamos formas en [C, H, W], donde C representa el número de canales y H, W representa la altura y el ancho respectivamente, esta función devuelve un tensor de imagen con cuadros delimitadores. El cuadro delimitador debe contener cuatro puntos en formato como (xmin, ymin, xmax, ymax), el punto superior izquierdo = (xmin, ymin) y el punto inferior derecho = (xmax, ymax). 

Sintaxis: torch.utils.draw_bounding_boxes (imagen, cuadro)

Parámetro:

• imagen: Tensor de forma (C x H x W) 
• cuadro: tamaño de los cuadros delimitadores en (xmin, ymin, xmax, ymax).

Retorno: Tensor de imagen de dtype uint8 con cuadros delimitadores trazados.

La siguiente imagen se utiliza para la demostración:

 

Ejemplo 1:

El siguiente ejemplo es para saber cómo cambiar y rellenar colores en cuadros delimitadores.

# 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()

Producción:

 

Ejemplo 2:

El siguiente ejemplo es para comprender cómo dibujar múltiples cuadros delimitadores en una imagen.

# 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()

Producción: