Node.js环境下如何从图片URL获取ImageData等对象用于TensorFlow图像分类

Node.js 环境没有浏览器内置的 HTMLImageElement、HTMLCanvasElement、ImageData 等 DOM API，你之前使用的 tf.browser.fromPixels 是浏览器专属API，在Node环境下无法直接使用，也不需要强行构造DOM类型对象，可以直接用TensorFlow.js针对Node.js的原生API实现需求。

第一步：替换依赖 #

将浏览器版本的@tensorflow/tfjs替换为Node原生版本，获得图片解码能力：

npm install @tensorflow/tfjs-node @tensorflow-models/mobilenet @tensorflow-models/knn-classifier axios

第二步：修改代码适配Node环境 #

我们可以直接拉取图片二进制数据，用tf.node.decodeImage直接解码为后续推理需要的张量，不需要中转成DOM类型对象，同时修正你代码中的拼写错误、异步方法漏加await等问题：

// 引入Node版tfjs
const tf = require("@tensorflow/tfjs-node");
const mobilenetModule = require("@tensorflow-models/mobilenet");
const knnClassifier = require("@tensorflow-models/knn-classifier");
const axios = require("axios");

// 修正之前的拼写错误 classifer => classifier
const classifier = knnClassifier.create();

// 封装通用方法：从图片URL直接获取可用于推理的张量
async function loadImageFromUrl(url) {
  // 拉取图片二进制数据
  const res = await axios.get(url, {
    responseType: "arraybuffer"
  });
  // 解码二进制为RGB三通道图片张量
  return tf.node.decodeImage(new Uint8Array(res.data), 3);
}

async function start() {
  const mobilenet = await mobilenetModule.load();

  // 加载训练集图片
  const img0 = await loadImageFromUrl("https://upload.wikimedia.org/wikipedia/commons/thumb/b/b6/Image_created_with_a_mobile_phone.png/1200px-Image_created_with_a_mobile_phone.png");
  const logits0 = mobilenet.infer(img0, true);
  classifier.addExample(logits0, 0);

  const img1 = await loadImageFromUrl("https://www.pixsy.com/wp-content/uploads/2021/04/ben-sweet-2LowviVHZ-E-unsplash-1.jpeg");
  const logits1 = mobilenet.infer(img1, true);
  classifier.addExample(logits1, 0);

  // 加载待检测图片并预测
  const x = await loadImageFromUrl("https://cryptotvnetwork.com/wp-content/uploads/2021/04/4ee1ad2ffbb00866fb7c55c61786e95d.jpg");
  const xlogits = mobilenet.infer(x, true);
  // 预测方法为异步，需加await
  const p = await classifier.predictClass(xlogits);
  console.log("预测结果：", p);

  // 手动释放张量内存，避免内存泄漏
  [img0, img1, x, logits0, logits1, xlogits].forEach(t => t?.dispose());
}

start().catch(err => console.error("运行错误：", err));

可选：获取ImageData对象 #

如果你有其他逻辑确实需要用到ImageData类型，可以安装canvas包获得Node环境下的ImageData构造能力：

npm install canvas

转换示例：

const { ImageData } = require("canvas");

// 假设imageTensor为已解码的图片张量
const pixelArray = await tf.browser.toPixels(imageTensor);
const imageData = new ImageData(pixelArray, imageTensor.shape[1], imageTensor.shape[0]);

内容的提问来源于stack exchange，提问作者Bob