如何在 PyTorch 中计算一组边界框的面积?
torchvision.io 包提供了执行不同 IO 操作的功能。为了计算边界框或一组边界框的面积,torchvision.io 包提供了 box_area() 函数。此函数将边界框作为输入参数,并返回每个框的面积。
边界框应为大小为 [N,4] 的 torch 张量,其中 N 是要计算面积的边界框的数量。每个边界框由坐标 (xmin, ymin, xmax, ymax) 指定。换句话说 - 0 ≤ xmin < xmax, 且 0 ≤ ymin < ymax. 计算出的面积是一个大小为 [N] 的 torch 张量。
要计算单个边界框的面积,我们将边界框张量进行 unsqueeze 操作,使其成为一个二维张量。
语法
torchvision.ops.box_area(boxes)
参数
boxes - 包含边界框的 [N,4] torch 张量。每个边界框都以 (xmin, ymin, xmax, ymax) 格式给出,其中 0 ≤ xmin < xmax, 且 0 ≤ ymin < ymax.
输出
它返回一个大小为 [N] 的 torch 张量,其中包含边界框的面积。
步骤
导入所需的库。在以下所有示例中,所需的 Python 库为 torch 和 torchvision。确保您已安装它们。
import torch import torchvision from torchvision.io import read_image from torchvision.utils import draw_bounding_boxes from torchvision.ops import box_area
使用 image_read() 函数读取 JPEG 或 PNG 图像。使用图像类型(.jpg 或 .png)指定完整的图像路径。此函数的输出是一个大小为 [image_channels, image_height, image_width] 的 torch 张量。
img = read_image('dog.png')
将边界框定义为 torch 张量。边界框张量的数据类型应为 torch.int。如果要计算单个边界框的面积,则对张量进行 unsqueeze 操作。
bbox = (310, 200, 485, 430) # convert the bbox to torch tensor bbox = torch.tensor(bbox, dtype=torch.int)
使用 box_area() 计算边界框的面积。可以选择将带有绘制边界框的图像分配给一个新的变量。
area = box_area(bbox)
使用 draw_bounding_boxes() 函数在图像上绘制边界框。我们将计算出的面积作为标签放在边界框内。
labels= [f"bbox area = {area.item()}"] img=draw_bounding_boxes(img, bbox, labels= labels, width=3,colors=(255,255,0))
将图像张量转换为 PIL 图像并显示它。
img = torchvision.transforms.ToPILImage()(img) img.show()
输入图像
我们将在以下示例中使用这些图像作为输入文件。
示例 1
在以下 Python 程序中,我们计算单个边界框的面积,并将此面积作为标签放在图像上,并显示图像。
# Import the required libraries import torch import torchvision from torchvision.io import read_image from torchvision.utils import draw_bounding_boxes from torchvision.ops import box_area # read the input image img = read_image('dog.png') # bounding box in (xmin, ymin, xmax, ymax) format # top-left point=(xmin, ymin), bottom-right point = (xmax, ymax)bbox = (310, 200, 485, 430) # convert the bbox to torch tensor bbox = torch.tensor(bbox, dtype=torch.int) print(bbox) print(bbox.size()) # unsqueeze the bbox to make it 2D bbox = bbox.unsqueeze(0) print(bbox.size()) # Compute the bounding box area area = box_area(bbox) print("BBOX area:", area) labels= [f"bbox area = {area.item()}"] img=draw_bounding_boxes(img, bbox, labels= labels, width=3,colors=(255,255,0)) # b=a.permute(1,2,0) # plt.imshow(b) # plt.show() img = torchvision.transforms.ToPILImage()(img) img.show()
输出
tensor([310, 200, 485, 430], dtype=torch.int32) torch.Size([4]) torch.Size([1, 4]) BBOX area: tensor([40250], dtype=torch.int32)
示例 2
在以下 Python 程序中,我们计算一组两个边界框的面积,并将这些面积作为标签放在图像上,并显示图像。
import torch from PIL import Image import torchvision from torchvision.io import read_image from torchvision.utils import draw_bounding_boxes from torchvision.ops import box_area img = read_image('catndog.png') # bounding box in (xmin, ymin, xmax, ymax) format bbox1 = [30, 45, 330, 450] bbox2 = [320, 150, 690, 460] bbox = [bbox1, bbox2] bbox = torch.tensor(bbox, dtype=torch.int) print(bbox) print(bbox.size()) area = box_area(bbox) labels = [f"bbox area ={a}" for a in area] print(labels) img=draw_bounding_boxes(img, bbox, labels = labels, width=3,colors=[(255,0,0),(0,255,0)]) img = torchvision.transforms.ToPILImage()(img) img.show()
输出
tensor([[ 30, 45, 330, 450], [320, 150, 690, 460]], dtype=torch.int32) torch.Size([2, 4]) ['bbox area =121500', 'bbox area =114700']