如何在Detectron2可视化结果中添加关键点名称标注？

方案一：手动遍历关键点并添加文本 #

你可以在调用draw_instance_predictions之后，手动提取每个实例的关键点坐标，结合Detectron2的draw_text函数，将关键点名称标注在对应位置。修改后的代码如下：

outputs = predictor(im[..., ::-1])
v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)

# 先绘制默认的预测结果（边界框、类别、关键点）
out = v.draw_instance_predictions(outputs["instances"].to("cpu"))

# 获取数据集元数据中的关键点名称
keypoint_names = MetadataCatalog.get(cfg.DATASETS.TRAIN[0]).keypoint_names
# 获取所有实例的关键点数据（形状为[N, K, 2]，N是实例数，K是关键点数量）
instances = outputs["instances"].to("cpu")
keypoints = instances.pred_keypoints

# 遍历每个实例的关键点，添加名称标注
for inst_keypoints in keypoints:
    for idx, (x, y) in enumerate(inst_keypoints):
        # 跳过未检测到的关键点（坐标为0）
        if x == 0 and y == 0:
            continue
        # 将文本标注在关键点右上方，避免遮挡关键点
        text_pos = (x + 10, y - 10)
        # 使用draw_text添加标注，设置字体大小和颜色
        out = v.draw_text(keypoint_names[idx], text_pos, font_size=16, color="white")

# 可视化结果
plt.figure(figsize=(60,30))
plt.imshow(out.get_image()[..., ::-1])
plt.show()

方案二：自定义Visualizer类（长期复用方案） #

如果需要每次可视化都自动添加关键点名称，可以自定义一个继承自Visualizer的类，重写draw_keypoints方法，在绘制关键点的同时添加名称：

from detectron2.utils.visualizer import Visualizer

class KeypointNamedVisualizer(Visualizer):
    def draw_keypoints(self, keypoints, colors=None, radius=3, use_normalized_coords=False):
        # 先调用父类方法绘制关键点
        super().draw_keypoints(keypoints, colors, radius, use_normalized_coords)
        # 获取关键点名称
        keypoint_names = self.metadata.keypoint_names
        # 遍历每个关键点添加名称
        for idx, (x, y) in enumerate(keypoints):
            if x == 0 and y == 0:
                continue
            text_pos = (x + 10, y - 10)
            self.draw_text(keypoint_names[idx], text_pos, font_size=16, color="white")
        return self.output

# 使用自定义的Visualizer
outputs = predictor(im[..., ::-1])
v = KeypointNamedVisualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
out = v.draw_instance_predictions(outputs["instances"].to("cpu"))

plt.figure(figsize=(60,30))
plt.imshow(out.get_image()[..., ::-1])
plt.show()

注意事项 #

• 确保你的数据集元数据中已经正确配置了keypoint_names（比如注册数据集时设置keypoint_names=["keypoint_up", "keypoint_down"]），才能正确获取关键点名称。
• 可根据图像尺寸和关键点分布调整text_pos的偏移量，避免文本遮挡关键点。
• font_size参数可配合Visualizer的scale值调整，保证文本大小适配图像。

内容的提问来源于stack exchange，提问作者Amit