docs/guide/42_Keypoint_Detection.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2f03ad73",
   "metadata": {},
   "source": [
    "# 关键点检测\n",
    "\n",
    "关键点检测用于识别和定位图像中的特定兴趣点。这些关键点也称为标志点，代表物体的有意义特征，例如面部特征或物体部分。这些模型输入一张图像，并返回以下输出：\n",
    "\n",
    "- **关键点和分数**：兴趣点及其置信度分数。\n",
    "- **描述符**：每个关键点周围图像区域的表示，捕捉其纹理、梯度、方向等属性。\n",
    "\n",
    "在本指南中，我们将展示如何从图像中提取关键点。\n",
    "\n",
    "在本教程中，我们将使用 [SuperPoint](./model_doc/superpoint.md)，这是一个用于关键点检测的基础模型。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a337055",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoImageProcessor, SuperPointForKeypointDetection\n",
    "\n",
    "# 初始化处理器和模型\n",
    "processor = AutoImageProcessor.from_pretrained(\"magic-leap-community/superpoint\")\n",
    "model = SuperPointForKeypointDetection.from_pretrained(\"magic-leap-community/superpoint\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6128891e",
   "metadata": {},
   "source": [
    "\n",
    "我们将在以下图像上测试模型。\n",
    "\n",
    "![蜜蜂](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/bee.jpg)\n",
    "![猫](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/cats.png)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff5bc27f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from PIL import Image\n",
    "import requests\n",
    "import cv2\n",
    "\n",
    "# 下载图像\n",
    "url_image_1 = \"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/bee.jpg\"\n",
    "image_1 = Image.open(requests.get(url_image_1, stream=True).raw)\n",
    "url_image_2 = \"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/cats.png\"\n",
    "image_2 = Image.open(requests.get(url_image_2, stream=True).raw)\n",
    "\n",
    "# 将图像放入列表中\n",
    "images = [image_1, image_2]\n",
    "\n",
    "# 处理输入并进行推理\n",
    "inputs = processor(images, return_tensors=\"pt\").to(model.device, model.dtype)\n",
    "outputs = model(**inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6db215a",
   "metadata": {},
   "source": [
    "\n",
    "模型输出包括每个图像中的相对关键点、描述符、掩码和分数。掩码突出显示图像中存在关键点的区域。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ac527c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 输出示例\n",
    "outputs.keypoints\n",
    "outputs.scores\n",
    "outputs.descriptors\n",
    "outputs.mask"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44ee8370",
   "metadata": {},
   "source": [
    "\n",
    "为了在图像中标绘实际的关键点，我们需要对输出进行后处理。为此，我们需要将实际的图像尺寸传递给 `post_process_keypoint_detection` 函数，同时传递输出。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2b86047",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 获取图像尺寸\n",
    "image_sizes = [(image.size[1], image.size[0]) for image in images]\n",
    "# 进行后处理\n",
    "outputs = processor.post_process_keypoint_detection(outputs, image_sizes)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45d74024",
   "metadata": {},
   "source": [
    "\n",
    "现在，输出是一个字典列表，每个字典包含处理后的关键点、分数和描述符。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9557a69e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 输出示例\n",
    "outputs[0]['keypoints']\n",
    "outputs[0]['scores']\n",
    "outputs[0]['descriptors']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62a4c8d2",
   "metadata": {},
   "source": [
    "\n",
    "我们可以使用这些数据来标绘关键点。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1397c10",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "\n",
    "# 遍历每个图像\n",
    "for i in range(len(images)):\n",
    "    keypoints = outputs[i][\"keypoints\"]\n",
    "    scores = outputs[i][\"scores\"]\n",
    "    descriptors = outputs[i][\"descriptors\"]\n",
    "\n",
    "    # 将张量转换为 NumPy 数组\n",
    "    keypoints = outputs[i][\"keypoints\"].detach().numpy()\n",
    "    scores = outputs[i][\"scores\"].detach().numpy()\n",
    "    image = images[i]\n",
    "    image_width, image_height = image.size\n",
    "\n",
    "    # 绘制图像和关键点\n",
    "    plt.axis('off')\n",
    "    plt.imshow(image)\n",
    "    plt.scatter(\n",
    "        keypoints[:, 0],\n",
    "        keypoints[:, 1],\n",
    "        s=scores * 100,  # 置信度分数影响点的大小\n",
    "        c='cyan',  # 点的颜色\n",
    "        alpha=0.4  # 透明度\n",
    "    )\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "557ec12d",
   "metadata": {},
   "source": [
    "\n",
    "下面是输出结果。\n",
    "\n",
    "![蜜蜂](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/bee_keypoint.png)\n",
    "![猫](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/cats_keypoint.png)\n"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}