from transformers import AutoImageProcessor, SuperPointForKeypointDetection

# 初始化处理器和模型
processor = AutoImageProcessor.from_pretrained("magic-leap-community/superpoint")
model = SuperPointForKeypointDetection.from_pretrained("magic-leap-community/superpoint")

import torch
from PIL import Image
import requests
import cv2

# 下载图像
url_image_1 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/bee.jpg"
image_1 = Image.open(requests.get(url_image_1, stream=True).raw)
url_image_2 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/cats.png"
image_2 = Image.open(requests.get(url_image_2, stream=True).raw)

# 将图像放入列表中
images = [image_1, image_2]

# 处理输入并进行推理
inputs = processor(images, return_tensors="pt").to(model.device, model.dtype)
outputs = model(**inputs)

# 输出示例
outputs.keypoints
outputs.scores
outputs.descriptors
outputs.mask

# 获取图像尺寸
image_sizes = [(image.size[1], image.size[0]) for image in images]
# 进行后处理
outputs = processor.post_process_keypoint_detection(outputs, image_sizes)

# 输出示例
outputs[0]['keypoints']
outputs[0]['scores']
outputs[0]['descriptors']

import matplotlib.pyplot as plt
import torch

# 遍历每个图像
for i in range(len(images)):
    keypoints = outputs[i]["keypoints"]
    scores = outputs[i]["scores"]
    descriptors = outputs[i]["descriptors"]

    # 将张量转换为 NumPy 数组
    keypoints = outputs[i]["keypoints"].detach().numpy()
    scores = outputs[i]["scores"].detach().numpy()
    image = images[i]
    image_width, image_height = image.size

    # 绘制图像和关键点
    plt.axis('off')
    plt.imshow(image)
    plt.scatter(
        keypoints[:, 0],
        keypoints[:, 1],
        s=scores * 100,  # 置信度分数影响点的大小
        c='cyan',  # 点的颜色
        alpha=0.4  # 透明度
    )
    plt.show()

关键点检测¶