Vitpose Base Simple

🚀 变换工具库（Transformers.js）

将nielsr/vitpose-base-simple模型适配Transformers.js，用于关键点检测

🚀 快速开始

本项目将 nielsr/vitpose-base-simple 模型的ONNX权重进行处理，以适配Transformers.js库，实现关键点检测功能。

📦 安装指南

如果你还未安装 Transformers.js JavaScript库，可以使用以下命令从 NPM 进行安装：

npm i @huggingface/transformers

💻 使用示例

基础用法

以下是使用 onnx-community/vitpose-base-simple 进行姿态估计的示例代码：

import { AutoModel, AutoImageProcessor, RawImage } from '@huggingface/transformers';

// 加载模型和处理器
const model_id = 'onnx-community/vitpose-base-simple';
const model = await AutoModel.from_pretrained(model_id);
const processor = await AutoImageProcessor.from_pretrained(model_id);

// 加载图像并准备输入
const url = 'https://huggingface.co/datasets/Xenova/transformers.js-docs/resolve/main/ryan-gosling.jpg';
const image = await RawImage.read(url);
const inputs = await processor(image);

// 预测热力图
const { heatmaps } = await model(inputs);

// 后处理热力图以获取关键点和分数
const boxes = [[[0, 0, image.width, image.height]]];
const results = processor.post_process_pose_estimation(heatmaps, boxes)[0][0];
console.log(results);

高级用法

可选地，你可以对输出结果进行可视化（以下为Node.js环境下的示例，使用了 canvas 库）：

import { createCanvas, createImageData } from 'canvas';

// 创建画布并绘制图像
const canvas = createCanvas(image.width, image.height);
const ctx = canvas.getContext('2d');
const imageData = createImageData(image.rgba().data, image.width, image.height);
ctx.putImageData(imageData, 0, 0);

// 在关键点之间绘制连线
const points = results.keypoints;
ctx.lineWidth = 4;
ctx.strokeStyle = 'blue';
for (const [i, j] of model.config.edges) {
    const [x1, y1] = points[i];
    const [x2, y2] = points[j];
    ctx.beginPath();
    ctx.moveTo(x1, y1);
    ctx.lineTo(x2, y2);
    ctx.stroke();
}

// 在每个关键点处绘制圆圈
ctx.fillStyle = 'red';
for (const [x, y] of points) {
    ctx.beginPath();
    ctx.arc(x, y, 8, 0, 2 * Math.PI);
    ctx.fill();
}

// 将图像保存到文件
import fs from 'fs';
const out = fs.createWriteStream('pose.png');
const stream = canvas.createPNGStream();
stream.pipe(out)
out.on('finish', () =>  console.log('The PNG file was created.'));

以下是输入图像和输出图像的对比：

输入图像	输出图像

⚠️ 重要提示

为ONNX权重单独创建一个仓库是一种临时解决方案，直到WebML获得更广泛的应用。如果你想让你的模型适用于Web环境，我们建议使用 🤗 Optimum 将模型转换为ONNX格式，并像本仓库一样组织你的仓库结构（将ONNX权重放在名为 onnx 的子文件夹中）。