TorchVision官方文档翻译为中文-示例库可视化实用程序-004

Posted 2022-07-14 wx62cecd4679aef

此示例演示了torchvision为可视化图像、边界框和分割遮罩提供的一些实用程序。

import torch
import numpy as np
import matplotlib.pyplot as plt

import torchvision.transforms.functional as F

#plt.rcParams[savefig.dpi] = 300      # 图片像素
#plt.rcParams["savefig.format"] = pdf # 保存为pdf LaTex排版使用
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=[])

可视化一个网格图像
make_grid（）函数可用于创建表示网格中多个图像的张量。此util需要一个数据类型为uint8的图像作为输入。

from torchvision.utils import make_grid
from torchvision.io import read_image
from pathlib import Path

dog1_int = read_image(str(Path(assets) / dog1.jpg))
dog2_int = read_image(str(Path(assets) / dog2.jpg))

grid = make_grid([dog1_int, dog2_int, dog1_int, dog2_int])
show(grid)

可视化目标框

我们可以使用draw_bounding_boxes（）在图像上绘制框。我们可以设置颜色、标签、宽度以及字体和字号。这些框采用（xmin、ymin、xmax、ymax）格式。

from torchvision.utils import draw_bounding_boxes


boxes = torch.tensor([[50, 50, 100, 200], [210, 150, 350, 430]], dtype=torch.float)
colors = ["blue", "yellow"]
result = draw_bounding_boxes(dog1_int, boxes, colors=colors, width=5)
show(result)

当然，我们也可以绘制由torchvision检测模型生成的边界框。下面是从fasterrcnn_resnet50_fpn()模型加载的更快的R-CNN模型的演示。您还可以尝试使用带有RetinaNet_resnet50_fpn（）的RetinaNet、带有ssdlite320_mobilenet_v3_large（）的SSDlite或带有ssd300_vgg16（）的SSD。

from torchvision.models.detection import fasterrcnn_resnet50_fpn
from torchvision.transforms.functional import convert_image_dtype


batch_int = torch.stack([dog1_int, dog2_int])
batch = convert_image_dtype(batch_int, dtype=torch.float)

model = fasterrcnn_resnet50_fpn(pretrained=True, progress=False)
model = model.eval()

outputs = model(batch)
print(outputs)

输出

Downloading: "https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth" to /home/matti/.cache/torch/hub/checkpoints/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth
[boxes: tensor([[215.9767, 171.1661, 402.0079, 378.7391],
        [344.6341, 172.6735, 357.6114, 220.1435],
        [153.1306, 185.5568, 172.9223, 254.7014]], grad_fn=<StackBackward>), labels: tensor([18,  1,  1]), scores: tensor([0.9989, 0.0701, 0.0611], grad_fn=<IndexBackward>), boxes: tensor([[ 23.5964, 132.4331, 449.9360, 493.0223],
        [225.8182, 124.6292, 467.2861, 492.2620],
        [ 18.5248, 135.4171, 420.9785, 479.2225]], grad_fn=<StackBackward>), labels: tensor([18, 18, 17]), scores: tensor([0.9980, 0.0879, 0.0671], grad_fn=<IndexBackward>)]

让我们绘制模型检测到的框。我们将只绘制得分大于给定阈值的框。

score_threshold = .8
dogs_with_boxes = [
    draw_bounding_boxes(dog_int, boxes=output[boxes][output[scores] > score_threshold], width=4)
    for dog_int, output in zip(batch_int, outputs)
]
show(dogs_with_boxes)

可视化分割任务

draw_segmentation_masks（）函数可用于在图像上绘制分割遮罩。语义分割和实例分割模型有不同的输出，因此我们将分别处理它们。

语义分割模型
我们将看到如何将其与torchvision的FCN Resnet-50一起使用，该FCN Resnet-50加载了FCN_resnet50（）。您还可以尝试使用DeepLabv3（DeepLabv3_resnet50（））或lraspp mobilenet模型（lraspp_mobilenet_v3_large（））。
让我们先看看模型的输出。请记住，一般来说，图像在传递到语义分割模型之前必须进行规范化。

from torchvision.models.segmentation import fcn_resnet50


model = fcn_resnet50(pretrained=True, progress=False)
model = model.eval()

normalized_batch = F.normalize(batch, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
output = model(normalized_batch)[out]
print(output.shape, output.min().item(), output.max().item())

输出

Downloading: "https://download.pytorch.org/models/fcn_resnet50_coco-1167a1af.pth" to /home/matti/.cache/torch/hub/checkpoints/fcn_resnet50_coco-1167a1af.pth
torch.Size([2, 21, 500, 500]) -7.089667320251465 14.858259201049805

如上所述，分割模型的输出是形状张量（batchsize、num类、H、W）。每个值都是非标准化分数，我们可以使用softmax将它们标准化为[0,1]。在softmax之后，我们可以将每个值解释为一个概率，该概率指示给定像素属于给定类别的可能性。

让我们绘制狗级和船级检测到的Mask：

sem_classes = [
    __background__, aeroplane, bicycle, bird, boat, bottle, bus,
    car, cat, chair, cow, diningtable, dog, horse, motorbike,
    person, pottedplant, sheep, sofa, train, tvmonitor
]
sem_class_to_idx = cls: idx for (idx, cls) in enumerate(sem_classes)

normalized_masks = torch.nn.functional.softmax(output, dim=1)

dog_and_boat_masks = [
    normalized_masks[img_idx, sem_class_to_idx[cls]]
    for img_idx in range(batch.shape[0])
    for cls in (dog, boat)
]

show(dog_and_boat_masks)

正如预期的那样，模特对狗的级别很有信心，但对船的级别没有那么多信心。

draw_segmentation_masks（）函数可用于在原始图像上绘制这些遮罩。此函数期望掩码为布尔掩码，但上面的掩码包含[0,1]中的概率。要获取布尔掩码，我们可以执行以下操作：

class_dim = 1
boolean_dog_masks = (normalized_masks.argmax(class_dim) == sem_class_to_idx[dog])
print(f"shape = boolean_dog_masks.shape, dtype = boolean_dog_masks.dtype")
show([m.float() for m in boolean_dog_masks])

上面我们定义布尔掩码的那一行有点神秘，但您可以将其理解为以下查询：“对于哪个像素，‘dog’最有可能是类？”

**注意：**当我们在这里使用规范化的掩码时，我们可以通过直接使用模型的非规范化分数得到相同的结果（因为softmax操作保留了顺序）。

现在我们有了布尔Mask，我们可以将它们与draw_segmentation_masks（）一起用于在原始图像上绘制它们：

from torchvision.utils import draw_segmentation_masks

dogs_with_masks = [
    draw_segmentation_masks(img, masks=mask, alpha=0.7)
    for img, mask in zip(batch_int, boolean_dog_masks)
]
show(dogs_with_masks)

我们可以为每幅图像绘制多个Mask！请记住，模型返回的掩码数量与类的数量相同。让我们询问与上面相同的查询，但这次是针对所有类，而不仅仅是狗类：“对于每个像素和每个C类，C类是最有可能的类？”
这一个有点复杂，所以我们将首先展示如何使用单个图像，然后我们将推广到批处理。

num_classes = normalized_masks.shape[1]
dog1_masks = normalized_masks[0]
class_dim = 0
dog1_all_classes_masks = dog1_masks.argmax(class_dim) == torch.arange(num_classes)[:, None, None]

print(f"dog1_masks shape = dog1_masks.shape, dtype = dog1_masks.dtype")
print(f"dog1_all_classes_masks = dog1_all_classes_masks.shape, dtype = dog1_all_classes_masks.dtype")

dog_with_all_masks = draw_segmentation_masks(dog1_int, masks=dog1_all_classes_masks, alpha=.6)
show(dog_with_all_masks)

输出

dog1_masks shape = torch.Size([21, 500, 500]), dtype = torch.float32
dog1_all_classes_masks = torch.Size([21, 500, 500]), dtype = torch.bool

我们可以在上图中看到，只画了两个面具：背景面具和狗面具。这是因为模型认为只有这两类是所有像素中最有可能的。如果模型检测到其他像素中最有可能的是另一个类，我们就会看到上面的遮罩。
删除背景掩码与传递masks=dog1_all_classes_masks[1:]一样简单，因为背景类是索引为0的类。
现在让我们做同样的事情，但对一整批图像。代码是类似的，但涉及到更多的维度。

class_dim = 1
all_classes_masks = normalized_masks.argmax(class_dim) == torch.arange(num_classes)[:, None, None, None]
print(f"shape = all_classes_masks.shape, dtype = all_classes_masks.dtype")
# The first dimension is the classes now, so we need to swap it
all_classes_masks = all_classes_masks.swapaxes(0, 1)

dogs_with_masks = [
    draw_segmentation_masks(img, masks=mask, alpha=.6)
    for img, mask in zip(batch_int, all_classes_masks)
]
show(dogs_with_masks)

输出

shape = torch.Size([21, 2, 500, 500]), dtype = torch.bool

实例分割模型
实例分割模型的输出与语义分割模型的输出显著不同。我们将在这里看到如何绘制此类模型的遮罩。让我们从分析MaskRCNN模型的输出开始。请注意，这些模型不需要对图像进行规范化，因此我们不需要使用规范化批处理。

我们将在这里描述MaskRCNN模型的输出。对象检测、实例分割和人物关键点检测中的模型都具有类似的输出格式，但其中一些模型可能具有额外的信息，如keypointrcnn_resnet50_fpn（），而其中一些模型可能没有掩码，如FasterrRcnn_resnet50_fpn（）

from torchvision.models.detection import maskrcnn_resnet50_fpn
model = maskrcnn_resnet50_fpn(pretrained=True, progress=False)
model = model.eval()

output = model(batch)
print(output)

输出

Downloading: "https://download.pytorch.org/models/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth" to /home/matti/.cache/torch/hub/checkpoints/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth
[boxes: tensor([[219.7444, 168.1722, 400.7378, 384.0263],
        [343.9716, 171.2287, 358.3447, 222.6263],
        [301.0303, 192.6917, 313.8879, 232.3154]], grad_fn=<StackBackward>), labels: tensor([18,  1,  1]), scores: tensor([0.9987, 0.7187, 0.6525], grad_fn=<IndexBackward>), masks: tensor([[[[0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.],
          ...,
          [0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.]]],


        [[[0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.],
          ...,
          [0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.]]],


        [[[0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.],
          ...,
          [0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.]]]], grad_fn=<UnsqueezeBackward0>), boxes: tensor([[ 44.6767, 137.9018, 446.5324, 487.3429],
        [  0.0000, 288.0053, 489.9293, 490.2352]], grad_fn=<StackBackward>), labels: tensor([18, 15]), scores: tensor([0.9978, 0.0697], grad_fn=<IndexBackward>), masks: tensor([[[[0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.],
          ...,
          [0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.]]],


        [[[0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.],
          ...,
          [0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.],
          [0., 0., 0.,  ..., 0., 0., 0.]]]], grad_fn=<UnsqueezeBackward0>)]

让我们来分析一下。对于批处理中的每个图像，模型输出一些检测（或实例）。每个输入图像的检测次数不同。每个实例都由其边界框、标签、分数和掩码描述。
输出的组织方式如下所示：输出是一个长度列表。列表中的每个条目对应于一个输入图像，它是一个带有“框”、“标签”、“分数”和“掩码”键的dict。与这些键关联的每个值都有num_instances元素。在我们上面的例子中，在第一个图像中检测到3个实例，在第二个图像中检测到2个实例。
框可以用上面的draw_bounding_boxes（）绘制，但这里我们更感兴趣的是遮罩。这些遮罩与我们在上面看到的语义分段模型的遮罩大不相同。

dog1_output = output[0]
dog1_masks = dog1_output[masks]
print(f"shape = dog1_masks.shape, dtype = dog1_masks.dtype, "
      f"min = dog1_masks.min(), max = dog1_masks.max()")

输出

shape = torch.Size([3, 1, 500, 500]), dtype = torch.float32, min = 0.0, max = 0.9999862909317017

这里，掩模对应于概率，该概率指示对于每个像素，其属于该实例的预测标签的可能性有多大。这些预测的标签对应于同一输出目录中的“labels”元素。让我们看看为第一个图像的实例预测了哪些标签。

inst_classes = [
    __background__, person, bicycle, car, motorcycle, airplane, bus,
    train, truck, boat, traffic light, fire hydrant, N/A, stop sign,
    parking meter, bench, bird, cat, dog, horse, sheep, cow,
    elephant, bear, zebra, giraffe, N/A, backpack, umbrella, N/A, N/A,
    handbag, tie, suitcase, frisbee, skis, snowboard, sports ball,
    kite, baseball bat, baseball glove, skateboard, surfboard, tennis racket,
    bottle, N/A, wine glass, cup, fork, knife, spoon, bowl,
    banana, apple, sandwich, orange, broccoli, carrot, hot dog, pizza,
    donut, cake, chair, couch, potted plant, bed, N/A, dining table,
    N/A, N/A, toilet, N/A, tv, laptop, mouse, remote, keyboard, cell phone,
    microwave, oven, toaster, sink, refrigerator, N/A, book,
    clock, vase, scissors, teddy bear, hair drier, toothbrush
]

inst_class_to_idx = cls: idx for (idx, cls) in enumerate(inst_classes)

print("For the first dog, the following instances were detected:")
print([inst_classes[label] for label in dog1_output[labels]])

输出

For the first dog, the following instances were detected:
[dog, person, person]

有趣的是，该模型在图像中检测到两个人。让我们继续画那些Mask。由于draw_segmentation_masks（）需要布尔掩码，因此我们需要将这些概率转换为布尔值。请记住，这些Mask的语义是“这个像素属于预测类的可能性有多大？”。因此，将这些掩码转换为布尔值的一种自然方法是以0.5的概率设置它们的阈值（也可以选择不同的阈值）。

proba_threshold = 0.5
dog1_bool_masks = dog1_output[masks] > proba_threshold
print(f"shape = dog1_bool_masks.shape, dtype = dog1_bool_masks.dtype")

# Theres an extra dimension (1) to the masks. We need to remove it
dog1_bool_masks = dog1_bool_masks.squeeze(1)

show(draw_segmentation_masks(dog1_int, dog1_bool_masks, alpha=0.9))

输出

shape = torch.Size([3, 1, 500, 500]), dtype = torch.bool

该模型似乎已经正确地检测到了狗，但它也把树和人混淆了。更仔细地观察分数将有助于我们绘制更相关的Mask：

print(dog1_output[scores])

输出

tensor([0.9987, 0.7187, 0.6525], grad_fn=<IndexBackward>)

显然，该模型对狗的检测不如对人的检测有信心。这是个好消息。在绘制遮罩时，我们只能要求那些得分良好的遮罩。让我们在这里使用0.75的分数阈值，并绘制第二只狗的Mask。

score_threshold = .75

boolean_masks = [
    out[masks][out[scores] > score_threshold] > proba_threshold
    for out in output
]

dogs_with_masks = [
    draw_segmentation_masks(img, mask.squeeze(1))
    for img, mask in zip(batch_int, boolean_masks)
]
show(dogs_with_masks)

第一幅图像中的两个“人物”Mask未被选中，因为它们的得分低于得分阈值。类似地，在第二幅图中，未选择类为15（对应于“bench”）的实例。更多计算机视觉与图形学相关资料，请关注微信公众号：计算机视觉与图形学实战

如果您认为上面的内容对您有一定的价值，可以对我们进行小小的赞助，来支持我们的工作，因为后续打算构建自己的网站，谢谢：