使用pytorch畫圖像標注框

使用draw_bounding_boxes可以方便的畫出圖像標注框

import torch
import numpy as np
import matplotlib.pyplot as plt
import torchvision.transforms.functional as F
from torchvision.io import read_image
from torchvision.utils import draw_bounding_boxes

plt.rcParams["savefig.bbox"] = 'tight'

def show(imgs):
    if not isinstance(imgs, list):
        imgs = [imgs]
    fix, axs = plt.subplots(ncols=len(imgs), squeeze=False)
    for i, img in enumerate(imgs):
        img = img.detach()
        img = F.to_pil_image(img)
        axs[0, i].imshow(np.asarray(img))
        axs[0, i].set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])

pic_1 = read_image('')

boxes = torch.tensor([[135, 37, 160, 198]], dtype=torch.float)
colors = ["blue"]
result = draw_bounding_boxes(pic_1, boxes, colors=colors, width=5)
show(result)