用 Next.js 和 Transformers.js 构建自己的 AI 工具

Oceansdeep
429 阅读3分钟

20241128-095843.gif

先上一张效果图，在浏览器中直接运行 AI 模型。

使用 Cursor 构建的 UI，试用链接:  build-your-own-rmbg-tool.vercel.app/。

Transformers.js 介绍

Transformers.js 是一个在浏览器中运行机器学习模型的开源库，无需服务端支持，专注于提供 Web 环境中的最先进的机器学习功能。其核心功能包括：

  • 跨模态任务支持：

    • 自然语言处理：文本分类、命名实体识别、问答、语言建模、摘要、翻译、多选题解答、文本生成等。
    • 计算机视觉：图像分类、目标检测、分割。
    • 音频处理：语音识别、音频分类、文本转语音。
    • 多模态任务：嵌入生成、零样本音频/图像分类与目标检测。

  • 与 Hugging Face Transformers 一致的 API：与 Python 版功能相当，支持预训练模型直接在浏览器中运行。

  • 基于 ONNX Runtime：依赖高效的 ONNX Runtime 来执行模型，兼容 PyTorch、TensorFlow 和 JAX 的模型，可通过 Optimum 工具轻松转换为 ONNX 格式。

通过 Transformers.js，开发者可在 Web 环境中实现高效的模型推理，支持多模态任务，适用于无无服务端的轻量级应用场景。

Transformers.js v3 支持了 WebGPU，下面是一些可用的模型：

e4317364-9159-498e-9cc6-8eaf3860d0dc.jpeg

构建一个 Next.js 应用

前置条件

在开始之前，请确保已安装以下工具：

  • Node.js：版本 18 或更高
  • npm：版本 9 或更高

第 1 步：初始化项目

首先，使用 create-next-app 创建一个新的 Next.js 应用程序，参考以下选项初始化项目：

√ What is your project named? ... build-your-own-rmbg-tool
√ Would you like to use TypeScript? ... Yes
√ Would you like to use ESLint? ... Yes
√ Would you like to use Tailwind CSS? ... Yes
√ Would you like to use `src/` directory? ... No 
√ Would you like to use App Router? (recommended) ... Yes
√ Would you like to customize the default import alias? ... Yes

第 2 步：安装并配置 Transformers.js

通过以下命令从 pnpm 安装 Transformers.js：

pnpm add @huggingface/transformers

接下来，修改 next.config.js 文件，在打包浏览器端代码时忽略 Node.js 专用模块：

import path from "path";  
import {fileURLToPath} from 'url';  
import {NextConfig} from "next";  
  
const __filename = fileURLToPath(import.meta.url);  
const __dirname = path.dirname(__filename);  
  
const nextConfig: NextConfig = {  
    output: 'export',  
    webpack: (config) => {  
        config.resolve.alias = {  
            ...config.resolve.alias,  
            "sharp$": false,  
            "onnxruntime-node$": false,  
            "@huggingface/transformers": path.resolve(__dirname, 'node_modules/@huggingface/transformers')  
        }  
        return config;  
    },  
};  
  
export default nextConfig;

第 3 步：使用 briaai / RMBG-1.4 模型处理图片

  1. 加载模型和处理器
const model = await AutoModel.from_pretrained('briaai/RMBG-1.4', {
    config: { model_type: 'custom' },
});

const processor = await AutoProcessor.from_pretrained('briaai/RMBG-1.4', {
    config: {
        do_normalize: true,
        do_pad: false,
        do_rescale: true,
        do_resize: true,
        image_mean: [0.5, 0.5, 0.5],
        feature_extractor_type: "ImageFeatureExtractor",
        image_std: [1, 1, 1],
        resample: 2,
        rescale_factor: 0.00392156862745098,
        size: { width: 1024, height: 1024 },
    }
});

2. 加载图片

const image = await RawImage.fromURL(imageUrl);

3. 获取蒙层数据

// 图片预处理
const { pixel_values } = await processor(image);

// 使用模型进行推理
const { output } = await model({ input: pixel_values });

// 调整蒙层大小
const mask = await RawImage.fromTensor(output[0].mul(255).to('uint8')).resize(image.width, image.height);

3. 用 Canvas 根据蒙层数据来绘制结果

// 创建一个新的 canvas 元素
const canvas = document.createElement('canvas');
canvas.width = image.width; // 设置 canvas 的宽度为原始图像宽度
canvas.height = image.height; // 设置 canvas 的高度为原始图像高度

// 获取 canvas 的 2D 绘图上下文
const context = canvas.getContext('2d');

// 在 canvas 上绘制原始图像
context.drawImage(image, 0, 0);

// 获取图像像素数据
const imagePixels = context.getImageData(0, 0, image.width, image.height);

// 遍历每个像素，更新其 alpha 通道（透明度）
for (let pixelIndex = 0; pixelIndex < mask.data.length; ++pixelIndex) {
    imagePixels.data[4 * pixelIndex + 3] = mask.data[pixelIndex]; // 将 mask 数据设置为 alpha 通道
}

// 将更新后的像素数据写回 canvas
context.putImageData(imagePixels, 0, 0);

// 导出图片
const resultImageUrl = canvas.toDataURL('image/png');

总结

完整代码: github.com/Oceansdeep7…

感谢阅读，希望本文对你有帮助，欢迎 Star 。如果有任何问题，欢迎在评论区进行交流。

avatar
文章
阅读
粉丝