如何在Node.js中为ChromaDB向量数据库嵌入图像向量

1. 本地使用预训练图像特征提取库 #

直接在Node.js环境中调用计算机视觉库提取图像特征向量，无需依赖外部API：

• TensorFlow.js + 预训练图像模型
  选用MobileNet、ResNet这类成熟的预训练模型，通过TensorFlow.js加载后，裁剪掉最终分类层即可得到特征向量。示例代码：

  const tf = require('@tensorflow/tfjs-node');
  const fs = require('fs');
  
  async function getImageVector(imagePath) {
    // 加载本地部署的预训练MobileNet模型（可从TensorFlow.js模型库下载文件）
    const model = await tf.loadLayersModel('./models/mobilenet_v1_0.25_224/model.json');
    // 构建特征提取器，移除最后两层分类相关层
    const featureExtractor = tf.model({
      inputs: model.inputs,
      outputs: model.layers[model.layers.length - 3].output
    });
  
    // 读取并预处理图像
    const imageBuffer = fs.readFileSync(imagePath);
    const imageTensor = tf.node.decodeJpeg(imageBuffer, 3)
      .resizeNearestNeighbor([224, 224])
      .toFloat()
      .expandDims();
  
    // 提取特征向量并转换为数组
    const features = featureExtractor.predict(imageTensor);
    const vector = await features.data();
    // 清理张量避免内存泄漏
    tf.dispose([model, featureExtractor, imageTensor, features]);
    return Array.from(vector);
  }
  

  提取的向量可直接存入ChromaDB，创建集合时需指定dimension参数匹配模型输出维度（如上述MobileNet配置输出1024维向量）。

2. 调用支持图像嵌入的云API #

若本地资源有限，可选择提供图像转向量能力的云服务，发送图像文件或URL即可获取向量：

const axios = require('axios');
const fs = require('fs');
const FormData = require('form-data');

async function getImageVectorViaAPI(imagePath) {
  const formData = new FormData();
  formData.append('image', fs.createReadStream(imagePath));
  const response = await axios.post('https://your-image-embedding-api-endpoint', formData, {
    headers: {
      'Authorization': 'Bearer YOUR_API_KEY',
      ...formData.getHeaders()
    }
  });
  return response.data.vector;
}

注意选择向量维度与ChromaDB兼容的服务（如CLIP模型输出512维向量）。

3. 本地部署图像嵌入模型（隐私优先方案） #

若需保障数据隐私，可通过Python部署CLIP、ViT这类图像文本匹配模型，再用Node.js调用本地接口：

• Python端FastAPI服务示例（仅核心逻辑）：

  from fastapi import FastAPI, File, UploadFile
  import clip
  import torch
  from PIL import Image
  
  app = FastAPI()
  device = "cuda" if torch.cuda.is_available() else "cpu"
  model, preprocess = clip.load("ViT-B/32", device=device)
  
  @app.post("/image-vector")
  async def generate_vector(file: UploadFile = File(...)):
    image = preprocess(Image.open(file.file)).unsqueeze(0).to(device)
    with torch.no_grad():
      image_features = model.encode_image(image)
    return {"vector": image_features.cpu().numpy().tolist()[0]}
  

• Node.js调用代码：

  const axios = require('axios');
  const fs = require('fs');
  const FormData = require('form-data');
  
  async function getLocalModelVector(imagePath) {
    const formData = new FormData();
    formData.append('file', fs.createReadStream(imagePath));
    const response = await axios.post('http://localhost:8000/image-vector', formData, {
      headers: formData.getHeaders()
    });
    return response.data.vector;
  }
  

关键注意事项 #

• 向量维度必须与ChromaDB集合的dimension配置完全匹配，否则无法存入或检索。
• 本地处理时需及时清理TensorFlow.js张量，避免内存泄漏。
• 云API调用需考虑成本、延迟及图像传输的带宽消耗。

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