pytorch画图怎样保存绘图结果 - 问答

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

0 赞

0 踩