TensorFlow.js实现AI换脸使用示例详解

node -v
npm -v

npm install @tensorflow/tfjs @tensorflow-models/face-landmarks-detection @tensorflow/tfjs-node
npm install canvas

// 导入所需的依赖包
const tf = require('@tensorflow/tfjs-node');
const faceLandmarksDetection = require('@tensorflow-models/face-landmarks-detection');
// 加载面部关键点检测模型
const loadModel = async () => {
  const model = await faceLandmarksDetection.load(
    faceLandmarksDetection.SupportedPackages.mediapipeFacemesh,
    { shouldLoadIrisModel: false }
  );
  return model;
};
// 调用loadModel函数加载模型
const model = await loadModel();

const { createCanvas, loadImage } = require('canvas');
// 加载两张图片
const loadImages = async () => {
  // 创建canvas并获取context
  const canvas = createCanvas(640, 480);
  const ctx = canvas.getContext('2d');
  // 加载sourceImage
  const sourceImage = await loadImage('source.jpg');
  canvas.width = sourceImage.width;
  canvas.height = sourceImage.height;
  ctx.drawImage(sourceImage, 0, 0);
  // 加载targetImage
  const targetImage = await loadImage('target.jpg');
  canvas.width = targetImage.width;
  canvas.height = targetImage.height;
  ctx.drawImage(targetImage, 0, 0);
  // 返回结果
  return { sourceImage, targetImage, canvas };
};
// 调用loadImages函数加载图片
const { sourceImage, targetImage, canvas } = await loadImages();

const extractFaceLandmarks = async (model, image) => {
  // 将image转换成tensor
  const tensor = tf.browser.fromPixels(image);
  // 使用estimateFaces方法得到关键点
  const predictions = await model.estimateFaces({
    input: tensor,
    returnTensors: false,
    flipHorizontal: false,
    predictIrises: false
  });
  // 释放tensor占用的内存
  tensor.dispose();
  return predictions;
};
// 提取sourceImage和targetImage的面部关键点
const sourceFaceLandmarks = await extractFaceLandmarks(model, sourceImage);
const targetFaceLandmarks = await extractFaceLandmarks(model, targetImage);

// 导入FaceMesh和Geometry2d类
const { FaceMesh } = require('@mediapipe/face_mesh');
const { Geometry2d } = require('@mediapipe/geometry');
// 定义applyWarp函数
const applyWarp = (sourceFaceLandmarks, targetFaceLandmarks, sourceImage, targetImage) => {
  // 创建FaceMesh、Geometry2d实例
  const faceMesh = new FaceMesh();
  const sourceGeometry = new Geometry2d();
  const targetGeometry = new Geometry2d();
  // 定义warpTriangles函数，使用Delaunay算法将source和target的关键点连接起来，得到一组三角形网格
  const warpTriangles = (sourcePoints, targetPoints) => {
    const delaunay = d3.Delaunay.from(sourcePoints);
    const { triangles } = delaunay;
    const sourceTriangles = triangles.map(i => sourcePoints[i]);
    const targetTriangles = triangles.map(i => targetPoints[i]);
    return { sourceTriangles, targetTriangles };
  };
  // 对source和target的关键点应用warpTriangles函数，得到sourceTriangles和targetTriangles
  const { sourceTriangles, targetTriangles } = warpTriangles(sourceFaceLandmarks[0].scaledMesh, targetFaceLandmarks[0].scaledMesh);
  // 将sourceTriangles和targetTriangles转换为Geometry2d实例
  sourceGeometry.setFromPoints(sourceTriangles.flat());
  targetGeometry.setFromPoints(targetTriangles.flat());
  // 计算仿射变换矩阵
  const warp = sourceGeometry.getAffineTransfORM(targetGeometry);
  // 创建canvas并绘制sourceImage
  const { canvas } = createCanvas(sourceImage.width, sourceImage.height);
  const ctx = canvas.getContext('2d');
  ctx.beginPath();
  ctx.moveTo(sourceFaceLandmarks[0].scaledMesh[0][0], sourceFaceLandmarks[0].scaledMesh[0][1]);
  for (let i = 1; i < sourceFaceLandmarks[0].scaledMesh.length; i++) {
    ctx.lineTo(sourceFaceLandmarks[0].scaledMesh[i][0], sourceFaceLandmarks[0].scaledMesh[i][1]);
  }
  ctx.closePath();
  ctx.clip();
  // 将targetImage按照仿射变换矩阵映射到sourceImage上
  ctx.setTransform(warp[0], warp[3], warp[1], warp[4], warp[2], warp[5]);
  ctx.drawImage(targetImage, 0, 0, targetImage.width, targetImage.height, 0, 0, sourceImage.width, sourceImage.height);
  // 将生成的图片转换为buffer并返回
  return canvas.toBuffer();
};
// 应用applyWarp函数
const result = applyWarp(sourceFaceLandmarks, targetFaceLandmarks, sourceImage, targetImage);

const fs = require('fs');
fs.writeFileSync('result.jpg', result);