快捷键

可视化工具

注意

试用 Colab转到结尾 下载完整的示例代码。

此示例说明了 torchvision 提供的一些用于可视化图像、边界框、分割掩码和关键点的工具。

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

import torchvision.transforms.functional as F


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


def show(imgs):
    if not isinstance(imgs, list):
        imgs = [imgs]
    fig, 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() 函数可用于创建表示网格中多个图像的张量。此工具需要 dtype 为 uint8 的单个图像作为输入。

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

dog1_int = decode_image(str(Path('../assets') / 'dog1.jpg'))
dog2_int = decode_image(str(Path('../assets') / 'dog2.jpg'))
dog_list = [dog1_int, dog2_int]

grid = make_grid(dog_list)
show(grid)
plot visualization utils

可视化边界框

我们可以使用 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)
plot visualization utils

当然,我们也可以绘制 torchvision 检测模型生成的边界框。以下是一个使用从 fasterrcnn_resnet50_fpn() 模型加载的 Faster R-CNN 模型进行演示。有关此类模型输出的更多详细信息,请参阅 实例分割模型

from torchvision.models.detection import fasterrcnn_resnet50_fpn, FasterRCNN_ResNet50_FPN_Weights


weights = FasterRCNN_ResNet50_FPN_Weights.DEFAULT
transforms = weights.transforms()

images = [transforms(d) for d in dog_list]

model = fasterrcnn_resnet50_fpn(weights=weights, progress=False)
model = model.eval()

outputs = model(images)
print(outputs)
[{'boxes': tensor([[215.9767, 171.1661, 402.0078, 378.7391],
        [344.6341, 172.6735, 357.6114, 220.1435],
        [153.1306, 185.5567, 172.9223, 254.7014]], grad_fn=<StackBackward0>), 'labels': tensor([18,  1,  1]), 'scores': tensor([0.9989, 0.0701, 0.0611], grad_fn=<IndexBackward0>)}, {'boxes': tensor([[ 23.5964, 132.4331, 449.9359, 493.0222],
        [225.8182, 124.6292, 467.2861, 492.2621],
        [ 18.5248, 135.4171, 420.9786, 479.2225]], grad_fn=<StackBackward0>), 'labels': tensor([18, 18, 17]), 'scores': tensor([0.9980, 0.0879, 0.0671], grad_fn=<IndexBackward0>)}]

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

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(dog_list, outputs)
]
show(dogs_with_boxes)
plot visualization utils

可视化分割掩码

draw_segmentation_masks() 函数可用于在图像上绘制分割掩码。语义分割和实例分割模型具有不同的输出,因此我们将分别处理它们。

语义分割模型

我们将看到如何使用 torchvision 的 FCN Resnet-50,使用 fcn_resnet50() 加载。让我们首先查看模型的输出。

from torchvision.models.segmentation import fcn_resnet50, FCN_ResNet50_Weights

weights = FCN_ResNet50_Weights.DEFAULT
transforms = weights.transforms(resize_size=None)

model = fcn_resnet50(weights=weights, progress=False)
model = model.eval()

batch = torch.stack([transforms(d) for d in dog_list])
output = model(batch)['out']
print(output.shape, output.min().item(), output.max().item())
Downloading: "https://download.pytorch.org/models/fcn_resnet50_coco-1167a1af.pth" to /root/.cache/torch/hub/checkpoints/fcn_resnet50_coco-1167a1af.pth
torch.Size([2, 21, 500, 500]) -7.089669704437256 14.858257293701172

如上所示,分割模型的输出是一个形状为 (batch_size, num_classes, H, W) 的张量。每个值都是一个非归一化分数,我们可以使用 softmax 将其归一化为 [0, 1]。在 softmax 之后,我们可以将每个值解释为概率,表示给定像素属于给定类别的可能性。

让我们绘制为狗类别和船类别检测到的掩码

sem_class_to_idx = {cls: idx for (idx, cls) in enumerate(weights.meta["categories"])}

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(len(dog_list))
    for cls in ('dog', 'boat')
]

show(dog_and_boat_masks)
plot visualization utils

正如预期的那样,模型对狗类别很有信心,但对船类别不太自信。

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])
plot visualization utils
shape = torch.Size([2, 500, 500]), dtype = torch.bool

我们在上面定义 boolean_dog_masks 的行有点难懂,但你可以将其理解为以下查询:“对于哪些像素来说,‘狗’ 是最有可能的类别?”

注意

虽然我们在这里使用的是 normalized_masks,但我们可以直接使用模型的非归一化分数获得相同的结果(因为 softmax 操作保留了顺序)。

现在我们有了布尔掩码,我们可以将它们与 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(dog_list, boolean_dog_masks)
]
show(dogs_with_masks)
plot visualization utils

我们可以为每个图像绘制多个掩码!请记住,模型返回的掩码数量与类别数量相同。让我们提出与上面相同的查询,但这次是针对所有类别,而不仅仅是狗类别:“对于每个像素和每个类别 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)
plot visualization utils
dog1_masks shape = torch.Size([21, 500, 500]), dtype = torch.float32
dog1_all_classes_masks = torch.Size([21, 500, 500]), dtype = torch.bool

我们可以在上面的图像中看到,只绘制了 2 个掩码:背景掩码和狗掩码。这是因为模型认为只有这两个类别在所有像素中是最有可能的。如果模型检测到另一个类别在其他像素中是最有可能的,我们将在上面看到它的掩码。

移除背景掩码就像传递 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(dog_list, all_classes_masks)
]
show(dogs_with_masks)
plot visualization utils
shape = torch.Size([21, 2, 500, 500]), dtype = torch.bool

实例分割模型

实例分割模型的输出与语义分割模型有很大不同。我们将在这里看到如何绘制此类模型的掩码。让我们首先分析 Mask-RCNN 模型的输出。请注意,这些模型不需要图像进行归一化,因此我们不需要使用归一化批处理。

注意

我们将在此描述 Mask-RCNN 模型的输出。中的模型 目标检测、实例分割和人物关键点检测 都具有类似的输出格式,但其中一些模型可能具有额外的信息,如 keypointrcnn_resnet50_fpn() 的关键点,而其中一些模型可能没有掩码,如 fasterrcnn_resnet50_fpn()

from torchvision.models.detection import maskrcnn_resnet50_fpn, MaskRCNN_ResNet50_FPN_Weights

weights = MaskRCNN_ResNet50_FPN_Weights.DEFAULT
transforms = weights.transforms()

images = [transforms(d) for d in dog_list]

model = maskrcnn_resnet50_fpn(weights=weights, progress=False)
model = model.eval()

output = model(images)
print(output)
Downloading: "https://download.pytorch.org/models/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth" to /root/.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=<StackBackward0>), 'labels': tensor([18,  1,  1]), 'scores': tensor([0.9987, 0.7187, 0.6525], grad_fn=<IndexBackward0>), '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.9292, 490.2352]], grad_fn=<StackBackward0>), 'labels': tensor([18, 15]), 'scores': tensor([0.9978, 0.0697], grad_fn=<IndexBackward0>), '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>)}]

让我们对其进行分解。对于批处理中的每张图像,模型都会输出一些检测结果(或实例)。每张输入图像的检测结果数量各不相同。每个实例都由其边界框、标签、得分和掩码来描述。

输出的组织方式如下:输出是一个长度为 batch_size 的列表。列表中的每个条目对应于一张输入图像,它是一个包含键‘boxes’、‘labels’、‘scores’ 和 ‘masks’ 的字典。与这些键关联的每个值都包含 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”元素。让我们看看第一张图像中实例的预测标签。

print("For the first dog, the following instances were detected:")
print([weights.meta["categories"][label] for label in dog1_output['labels']])
For the first dog, the following instances were detected:
['dog', 'person', 'person']

有趣的是,模型检测到图像中有两个人。让我们继续绘制这些掩码。由于 draw_segmentation_masks() 需要布尔掩码,我们需要将这些概率转换为布尔值。请记住,这些掩码的语义是“这个像素属于预测类的可能性有多大?”。因此,将这些掩码转换为布尔值的自然方法是使用 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}")

# There's 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))
plot visualization utils
shape = torch.Size([3, 1, 500, 500]), dtype = torch.bool

该模型似乎已正确检测到狗,但它也将树木误认为人。更仔细地观察分数将有助于我们绘制更相关的掩码。

print(dog1_output['scores'])
tensor([0.9987, 0.7187, 0.6525], grad_fn=<IndexBackward0>)

很明显,模型对狗检测的置信度高于对人检测的置信度。这是好消息。在绘制掩码时,我们可以只要求那些得分较高的掩码。让我们在这里使用 0.75 的分数阈值,并绘制第二只狗的掩码。

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(dog_list, boolean_masks)
]
show(dogs_with_masks)
plot visualization utils

第一张图像中的两个“人”掩码没有被选中,因为它们的分数低于分数阈值。同样,在第二张图像中,类别为 15(对应于“长凳”)的实例没有被选中。

可视化关键点

可以使用 draw_keypoints() 函数在图像上绘制关键点。我们将看到如何使用它与使用 keypointrcnn_resnet50_fpn() 加载的 torchvision 的 KeypointRCNN 一起使用。首先,我们来看看模型的输出。

from torchvision.models.detection import keypointrcnn_resnet50_fpn, KeypointRCNN_ResNet50_FPN_Weights
from torchvision.io import decode_image

person_int = decode_image(str(Path("../assets") / "person1.jpg"))

weights = KeypointRCNN_ResNet50_FPN_Weights.DEFAULT
transforms = weights.transforms()

person_float = transforms(person_int)

model = keypointrcnn_resnet50_fpn(weights=weights, progress=False)
model = model.eval()

outputs = model([person_float])
print(outputs)
Downloading: "https://download.pytorch.org/models/keypointrcnn_resnet50_fpn_coco-fc266e95.pth" to /root/.cache/torch/hub/checkpoints/keypointrcnn_resnet50_fpn_coco-fc266e95.pth
[{'boxes': tensor([[124.3751, 177.9242, 327.6354, 574.7064],
        [124.3625, 180.7574, 290.1061, 390.7958]], grad_fn=<StackBackward0>), 'labels': tensor([1, 1]), 'scores': tensor([0.9998, 0.1070], grad_fn=<IndexBackward0>), 'keypoints': tensor([[[208.0176, 214.2408,   1.0000],
         [208.0176, 207.0375,   1.0000],
         [197.8246, 210.6392,   1.0000],
         [208.0176, 211.8398,   1.0000],
         [178.6378, 217.8425,   1.0000],
         [221.2086, 253.8590,   1.0000],
         [160.6502, 269.4662,   1.0000],
         [243.9929, 304.2822,   1.0000],
         [138.4655, 328.8935,   1.0000],
         [277.5698, 340.8990,   1.0000],
         [153.4551, 374.5144,   1.0000],
         [226.0053, 375.7150,   1.0000],
         [226.0053, 370.3125,   1.0000],
         [221.8082, 455.5516,   1.0000],
         [273.9723, 448.9486,   1.0000],
         [193.6275, 546.1932,   1.0000],
         [273.3727, 545.5930,   1.0000]],

        [[207.8327, 214.6636,   1.0000],
         [207.2343, 207.4622,   1.0000],
         [198.2590, 209.8627,   1.0000],
         [208.4310, 210.4628,   1.0000],
         [178.5134, 218.2642,   1.0000],
         [219.7997, 251.8704,   1.0000],
         [162.3579, 269.2736,   1.0000],
         [245.5288, 304.6800,   1.0000],
         [138.4238, 330.4848,   1.0000],
         [278.4382, 346.0876,   1.0000],
         [153.3826, 374.8929,   1.0000],
         [233.5618, 368.2917,   1.0000],
         [225.7832, 367.6916,   1.0000],
         [289.8069, 357.4897,   1.0000],
         [245.5288, 389.8956,   1.0000],
         [281.4300, 349.0882,   1.0000],
         [209.0294, 389.8956,   1.0000]]], grad_fn=<CopySlices>), 'keypoints_scores': tensor([[16.0163, 16.6672, 15.8312,  4.6510, 14.2053,  8.8280,  9.1136, 12.2084,
         12.1901, 13.8453, 10.7090,  5.5852,  7.5005, 11.3378,  9.3700,  8.2987,
          8.4479],
        [12.9326, 13.8158, 14.9053,  3.9368, 12.9585,  6.4240,  6.8328, 10.4227,
          9.2907, 10.1066, 10.1019,  0.1822,  4.3058, -4.9904, -2.7409, -2.7874,
         -3.9329]], grad_fn=<CopySlices>)}]

正如我们所见,输出包含一个字典列表。输出列表的长度为 batch_size。我们目前只有一张图像,所以列表的长度为 1。列表中的每个条目对应于一个输入图像,它是一个包含键 boxeslabelsscoreskeypointskeypoint_scores 的字典。与这些键关联的每个值都包含 num_instances 个元素。在上面的例子中,图像中检测到 2 个实例。

kpts = outputs[0]['keypoints']
scores = outputs[0]['scores']

print(kpts)
print(scores)
tensor([[[208.0176, 214.2408,   1.0000],
         [208.0176, 207.0375,   1.0000],
         [197.8246, 210.6392,   1.0000],
         [208.0176, 211.8398,   1.0000],
         [178.6378, 217.8425,   1.0000],
         [221.2086, 253.8590,   1.0000],
         [160.6502, 269.4662,   1.0000],
         [243.9929, 304.2822,   1.0000],
         [138.4655, 328.8935,   1.0000],
         [277.5698, 340.8990,   1.0000],
         [153.4551, 374.5144,   1.0000],
         [226.0053, 375.7150,   1.0000],
         [226.0053, 370.3125,   1.0000],
         [221.8082, 455.5516,   1.0000],
         [273.9723, 448.9486,   1.0000],
         [193.6275, 546.1932,   1.0000],
         [273.3727, 545.5930,   1.0000]],

        [[207.8327, 214.6636,   1.0000],
         [207.2343, 207.4622,   1.0000],
         [198.2590, 209.8627,   1.0000],
         [208.4310, 210.4628,   1.0000],
         [178.5134, 218.2642,   1.0000],
         [219.7997, 251.8704,   1.0000],
         [162.3579, 269.2736,   1.0000],
         [245.5288, 304.6800,   1.0000],
         [138.4238, 330.4848,   1.0000],
         [278.4382, 346.0876,   1.0000],
         [153.3826, 374.8929,   1.0000],
         [233.5618, 368.2917,   1.0000],
         [225.7832, 367.6916,   1.0000],
         [289.8069, 357.4897,   1.0000],
         [245.5288, 389.8956,   1.0000],
         [281.4300, 349.0882,   1.0000],
         [209.0294, 389.8956,   1.0000]]], grad_fn=<CopySlices>)
tensor([0.9998, 0.1070], grad_fn=<IndexBackward0>)

KeypointRCNN 模型检测到图像中有两个实例。如果你使用 draw_bounding_boxes() 绘制框,你会发现它们是人和冲浪板。如果我们查看分数,我们会发现模型对人的置信度远高于冲浪板。现在,我们可以设置一个置信度阈值,并绘制我们足够自信的实例。让我们设置 0.75 的阈值并过滤掉与人相对应的关键点。

tensor([[[208.0176, 214.2408,   1.0000],
         [208.0176, 207.0375,   1.0000],
         [197.8246, 210.6392,   1.0000],
         [208.0176, 211.8398,   1.0000],
         [178.6378, 217.8425,   1.0000],
         [221.2086, 253.8590,   1.0000],
         [160.6502, 269.4662,   1.0000],
         [243.9929, 304.2822,   1.0000],
         [138.4655, 328.8935,   1.0000],
         [277.5698, 340.8990,   1.0000],
         [153.4551, 374.5144,   1.0000],
         [226.0053, 375.7150,   1.0000],
         [226.0053, 370.3125,   1.0000],
         [221.8082, 455.5516,   1.0000],
         [273.9723, 448.9486,   1.0000],
         [193.6275, 546.1932,   1.0000],
         [273.3727, 545.5930,   1.0000]]], grad_fn=<IndexBackward0>)

很好,现在我们有了与人相对应的关键点。每个关键点由 x、y 坐标和可见性表示。现在,我们可以使用 draw_keypoints() 函数绘制关键点。请注意,实用程序需要 uint8 图像。

from torchvision.utils import draw_keypoints

res = draw_keypoints(person_int, keypoints, colors="blue", radius=3)
show(res)
plot visualization utils

正如我们所见,关键点在图像上显示为彩色圆圈。人的 coco 关键点是有序的,表示以下列表。

coco_keypoints = [
    "nose", "left_eye", "right_eye", "left_ear", "right_ear",
    "left_shoulder", "right_shoulder", "left_elbow", "right_elbow",
    "left_wrist", "right_wrist", "left_hip", "right_hip",
    "left_knee", "right_knee", "left_ankle", "right_ankle",
]

如果我们想连接关键点呢?这在创建姿势检测或动作识别时特别有用。我们可以使用 connectivity 参数轻松地连接关键点。仔细观察会发现,我们需要按以下顺序连接点才能构建人体骨架。

鼻子 -> 左眼 -> 左耳。 (0, 1), (1, 3)

鼻子 -> 右眼 -> 右耳。 (0, 2), (2, 4)

鼻子 -> 左肩 -> 左肘 -> 左腕。 (0, 5), (5, 7), (7, 9)

鼻子 -> 右肩 -> 右肘 -> 右腕。 (0, 6), (6, 8), (8, 10)

左肩 -> 左髋 -> 左膝 -> 左踝。 (5, 11), (11, 13), (13, 15)

右肩 -> 右髋 -> 右膝 -> 右踝。 (6, 12), (12, 14), (14, 16)

我们将创建一个包含这些关键点 ID 的列表,以便连接它们。

connect_skeleton = [
    (0, 1), (0, 2), (1, 3), (2, 4), (0, 5), (0, 6), (5, 7), (6, 8),
    (7, 9), (8, 10), (5, 11), (6, 12), (11, 13), (12, 14), (13, 15), (14, 16)
]

我们将上述列表传递给 connectivity 参数以连接关键点。

res = draw_keypoints(person_int, keypoints, connectivity=connect_skeleton, colors="blue", radius=4, width=3)
show(res)
plot visualization utils

看起来不错。

使用可见性绘制关键点

让我们看看另一个关键点预测模块产生的结果,并显示连接性。

prediction = torch.tensor(
    [[[208.0176, 214.2409, 1.0000],
      [000.0000, 000.0000, 0.0000],
      [197.8246, 210.6392, 1.0000],
      [000.0000, 000.0000, 0.0000],
      [178.6378, 217.8425, 1.0000],
      [221.2086, 253.8591, 1.0000],
      [160.6502, 269.4662, 1.0000],
      [243.9929, 304.2822, 1.0000],
      [138.4654, 328.8935, 1.0000],
      [277.5698, 340.8990, 1.0000],
      [153.4551, 374.5145, 1.0000],
      [000.0000, 000.0000, 0.0000],
      [226.0053, 370.3125, 1.0000],
      [221.8081, 455.5516, 1.0000],
      [273.9723, 448.9486, 1.0000],
      [193.6275, 546.1933, 1.0000],
      [273.3727, 545.5930, 1.0000]]]
)

res = draw_keypoints(person_int, prediction, connectivity=connect_skeleton, colors="blue", radius=4, width=3)
show(res)
plot visualization utils

发生了什么?预测新关键点的模型无法检测到滑板者左上半身隐藏的三个点。更准确地说,模型预测 (x, y, vis) = (0, 0, 0) 用于左眼、左耳和左髋。因此,我们肯定不想显示这些关键点和连接,你也不必。查看 draw_keypoints() 的参数,我们可以看到,我们可以将可见性张量作为附加参数传递。鉴于模型的预测,我们有可见性作为第三个关键点维度,我们只需要将其提取出来。让我们将 prediction 拆分为关键点坐标及其各自的可见性,并将两者都作为参数传递给 draw_keypoints()

coordinates, visibility = prediction.split([2, 1], dim=-1)
visibility = visibility.bool()

res = draw_keypoints(
    person_int, coordinates, visibility=visibility, connectivity=connect_skeleton, colors="blue", radius=4, width=3
)
show(res)
plot visualization utils

我们可以看到,未检测到的关键点没有被绘制,并且不可见的关键点连接被跳过。这可以减少包含多个检测的图像上的噪声,或者在我们这种情况下,当关键点预测模型错过了一些检测时。大多数 torch 关键点预测模型会为每个预测返回可见性,以便你使用它。我们之前使用的 keypointrcnn_resnet50_fpn() 模型也是如此。

脚本的总运行时间: (0 分钟 13.710 秒)

由 Sphinx-Gallery 生成的图库

文档

访问 PyTorch 的全面开发人员文档

查看文档

教程

获取面向初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源