在字典列表中显示带有其他元素的图像

时间:2019-06-13 20:51:16

标签: pytorch tensor

我正在用我自己的图像实施pytorch网站提供的预训练检测和分割。当我打印“输出”时,它给出了字典列表,其中包含框,遮罩,标签等信息。我想在一幅图像中显示所有这些信息,但是我只知道显示遮罩。

请访问此网站here

import torchvision
import torchvision.transforms as transforms
import tensorflow as tf
import matplotlib.pyplot as plt
import numpy as np
import PIL
from PIL import Image

model= torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)
model.eval()

image = PIL.Image.open('ClothingAttributeDataset/images/000020.jpg')
image_tensor = torchvision.transforms.functional.to_tensor(image)

output = model([image_tensor])
print(output) # it will give a output of list of dictionary

for i in range(14):
    img = output[0]["masks"]
    img = img.data.cpu().numpy()[i]
    img = np.transpose(img, (1, 2, 0)).reshape(750, 500)
    plt.imsave('objects/img%d' %i, img)

0 个答案:

没有答案