在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
函数将绘制的图像保存到指定的路径。