TensorBoard标注预测结果框在图片上

TensorBoard官方教程

使用以下命令之一安装tensorboard

pip install tensorboard
conda install tensorboard

一、函数

标注预测结果框在图片上

from matplotlib.pyplot import Line2D
from torch.utils.tensorboard import SummaryWriter
import matplotlib.pyplot as plt
import numpy as np

import torch
import torchvision
import torchvision.transforms as transforms

import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

def matplotlib_imshow(img, one_channel=False):
    if one_channel:
        img = img.mean(dim=0)
    img = img / 2 + 0.5     # unnormalize
    npimg = img.numpy()
    if one_channel:
        plt.imshow(npimg, cmap="Greys")
    else:
        plt.imshow(np.transpose(npimg, (1, 2, 0)))

def plot_classes_preds(images, pred):
    '''
        批量绘制图像以及预测标签
        pred (list): 
        pred[i]:(torch.Tensor)    
                (N, (x1, y1, x2, y2, conf, cls))
    '''
    fig = plt.figure()
    num_of_images = images.size(0)
    for i in np.arange(num_of_images):
        ax = fig.add_subplot(1, num_of_images, i+1, xticks=[], yticks=[])
        num_of_boxes = pred[i].size(0)
        for j in np.arange(num_of_boxes):
            x1, y1, x2, y2, conf, labels = pred[i][j]
            ax.add_line(Line2D((x1, x2), (y1, y1)))
            ax.add_line(Line2D((x1, x2), (y2, y2)))
            ax.add_line(Line2D((x1, x1), (y1, y2)))
            ax.add_line(Line2D((x2, x2), (y1, y2)))
            ax.text(x1,y2, "label: {0}\nconf:{1}".format(labels, conf),fontsize=8,color = 'r')
        matplotlib_imshow(images[i], one_channel=True)   
    plt.savefig('new.jpg')
    return fig

def log_images(images, pred):
    # default `log_dir` is "runs" - we'll be more specific here
    writer = SummaryWriter('runs/fashion_mnist_experiment_1')
    fig = plot_classes_preds(images, pred)
    writer.add_figure('predictions on images', fig)

运行方法，运行以上代码的python函数之后，在命令行运行以下代码

tensorboard --logdir runs

之后tensorboard界面将显示在浏览器上

二、样例

def create_net():
    # transforms
    transform = transforms.Compose(
        [transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,))])

    # datasets
    trainset = torchvision.datasets.FashionMNIST('./data',
                                                download=True,
                                                train=True,
                                                transform=transform)

    # dataloaders
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
                                            shuffle=True, num_workers=2)

    # get some random training images
    dataiter = iter(trainloader)
    images, labels = dataiter.next()
    return images, labels

if __name__ == '__main__':
    images, labels = create_net()
    # pred = torch.tensor([[0, 0, 24, 24, 1, 0],[0, 0, 24, 24, 1, 0],[0, 0, 24, 24, 1, 0],[0, 0, 24, 24, 1, 0]])
    pred = []
    for i in range(4):
        # pred.append(torch.tensor([[6, 6, 12, 12, 0, 1]]))
        pred.append(torch.tensor([[6, 6, 12, 12, 0, i], [0, 0, 24, 24, 0, 2*i]]))
    # print(pred)
    log_images(images, pred)

运行之后在浏览器打开，显示界面如下：

放大后可以看到标注的框和标签、置信度