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 的子文件夾中）。