在PyTorch中，您可以使用torchvision.utils.draw_bounding_boxes函数来绘制边界框并在保存图像时显示它们。以下是一个示例代码，展示了如何使用此函数绘制边界框并保存图像：
import torch
from torchvision import transforms, utils
from PIL import Image
import os
# 加载图像并进行预处理
image = Image.open("path/to/your/image.jpg")
preprocess = transforms.Compose([
transforms.Resize((256, 256)),
transforms.ToTensor(),
])
image = preprocess(image).unsqueeze(0)
# 定义边界框及其类别
bboxes = torch.tensor([
[50, 50, 100, 100], # x1, y1, x2, y2
[150, 150, 200, 200],
])
labels = torch.tensor([0, 1]) # 假设0和1是类别索引
# 绘制边界框
draw = utils.draw_bounding_boxes(image, bboxes, labels)
# 保存图像
output_path = "path/to/save/image_with_bboxes.jpg"
utils.save_image(draw, output_path)
在这个示例中，我们首先加载图像并使用torchvision.transforms对其进行预处理。然后，我们定义了边界框及其类别，并使用torchvision.utils.draw_bounding_boxes函数在图像上绘制这些边界框。最后，我们使用utils.save_image函数将绘制的图像保存到指定的路径。
