視覺化工具

注意

在 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)

視覺化邊界框

我們可以使用 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() 模型載入的 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)

視覺化分割遮罩

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)

正如預期的那樣，該模型對狗類別充滿信心，但對船類別則不太有信心。

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])

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)

我們可以為每個圖像繪製多個遮罩！請記住，模型返回的遮罩與類別一樣多。讓我們提出與上面相同的查詢，但這次針對*所有*類別，而不僅僅是狗類別：“對於每個像素和每個類別 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

我們可以在上面的圖像中看到僅繪製了 2 個遮罩：背景遮罩和狗的遮罩。這是因為模型認為只有這 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)

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' 的 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

此處的遮罩對應於機率，指示對於每個像素，它屬於該實例的預測標籤的可能性有多大。這些預測的標籤對應於同一輸出 dict 中的 '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))

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

該模型似乎已正確檢測到狗，但它也將樹木與人混淆了。更仔細地查看分數將有助於我們繪製更多相關的遮罩

print(dog1_output['scores'])

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

顯然，該模型對狗的檢測比對人的檢測更有信心。這是個好消息。在繪製遮罩時，我們可以要求只繪製那些具有良好分數的遮罩。讓我們在此處使用 .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)

第一張圖片中的兩個「人」的遮罩未被選取，因為它們的分數低於分數閾值。 類似地，在第二張圖片中，類別為 15（對應於「長椅」）的實例未被選取。

視覺化關鍵點

draw_keypoints() 函數可用於在影像上繪製關鍵點。 我們將會看到如何搭配 torchvision 的 KeypointRCNN 使用，該模型是透過 keypointrcnn_resnet50_fpn() 載入的。 我們將首先查看模型的輸出。

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.4654, 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.5289, 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.5289, 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.3057, -4.9904, -2.7409, -2.7874,
         -3.9329]], grad_fn=<CopySlices>)}]

如我們所見，輸出包含字典的列表。 輸出列表的長度為 batch_size。 我們目前只有單一影像，因此列表的長度為 1。 列表中的每個條目都對應於一個輸入影像，它是一個具有 boxes、labels、scores、keypoints 和 keypoint_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.4654, 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.5289, 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.5289, 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 的閾值，並過濾掉對應於人的關鍵點。

detect_threshold = 0.75
idx = torch.where(scores > detect_threshold)
keypoints = kpts[idx]

print(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.4654, 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)

如我們所見，關鍵點顯示為影像上的彩色圓圈。 一個人的 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 參數輕鬆連接關鍵點。 仔細觀察會發現，我們需要按照以下順序連接這些點來建構人體骨骼。

nose -> left_eye -> left_ear. (0, 1), (1, 3)

nose -> right_eye -> right_ear. (0, 2), (2, 4)

nose -> left_shoulder -> left_elbow -> left_wrist. (0, 5), (5, 7), (7, 9)

nose -> right_shoulder -> right_elbow -> right_wrist. (0, 6), (6, 8), (8, 10)

left_shoulder -> left_hip -> left_knee -> left_ankle. (5, 11), (11, 13), (13, 15)

right_shoulder -> right_hip -> right_knee -> right_ankle. (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)

看起來還不錯。

繪製具有可見度的關鍵點

讓我們看看結果，另一個關鍵點預測模組產生了結果，並顯示了連線。

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)

那裡發生了什麼？ 預測新關鍵點的模型無法檢測到滑板運動員左上身上隱藏的三個點。 更準確地說，該模型預測左眼、左耳和左髖的 (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)

我們可以發現未檢測到的關鍵點沒有被繪製，並且跳過了不可見的關鍵點連線。 這可以減少具有多個檢測的影像上的雜訊，或者在像我們這樣的情況下，當關鍵點預測模型遺漏了一些檢測時。 大多數 torch 關鍵點預測模型都會傳回每個預測的可見度，隨時供你使用。 我們在第一個例子中使用的 keypointrcnn_resnet50_fpn() 模型也是如此。