在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
函數將繪制的圖像保存到指定的路徑。